"""Streaming TTS with clause-level chunking for Discord voice mode.

A worker thread consumes text -> produces 20ms PCM frames on a queue.Queue.
``EchoStreamingAudioSource`` pulls frames into Discord's audio thread so a
single ``voice_client.play()`` call lasts the whole turn (eliminates the
RTP gap between successive ``play()`` calls and the race with barge-in
``stop()``). See plan: src/voice/tts_stream.py (Pas 6 / Lane TTS),
Engineering decisions #6, #8, #15.
"""
from __future__ import annotations

import io
import logging
import queue
import re
import subprocess
import threading
import wave
from pathlib import Path
from typing import Iterator, List, Optional

import discord

from src.voice.normalize import normalize_for_tts
from tools.tts import synthesize


log = logging.getLogger(__name__)


# Discord wants 20ms of 16-bit 48kHz stereo PCM per frame.
# 48000 Hz * 0.020 s * 2 channels * 2 bytes = 3840 bytes.
FRAME_BYTES = 3840
TARGET_SAMPLE_RATE = 48000
TARGET_CHANNELS = 2
TARGET_SAMPLE_WIDTH = 2

# Sentinel pushed onto the text queue to ask the worker to exit cleanly.
_POISON = object()


# ---------- Clause segmentation ----------

# Split at Romanian sentence punctuation followed by whitespace. The
# trailing whitespace requirement protects mid-number (1.000), mid-decimal
# (12.5), and mid-abbreviation (M.D.) tokens, since none of those have a
# space right after the inner punctuation.
_CLAUSE_SPLIT = re.compile(r'(?<=[,;:.!?])\s+')


def clause_segments(text: str, min_words: int = 8) -> Iterator[str]:
    """Yield text in clause-sized chunks for streaming TTS.

    Splits at ``, ; : . ! ?`` boundaries (only when the punctuation is
    followed by whitespace, so numbers / decimals / abbreviations stay
    intact). Short clauses are buffered and merged with the next one
    until the accumulated chunk has at least ``min_words`` words. The
    final remainder is always yielded, even if it's shorter than
    ``min_words`` -- otherwise the tail of the response would never
    reach the TTS.
    """
    if text is None:
        return
    text = text.strip()
    if not text:
        return
    pieces = [p.strip() for p in _CLAUSE_SPLIT.split(text) if p and p.strip()]
    if not pieces:
        return
    buffer = ''
    for clause in pieces:
        buffer = (buffer + ' ' + clause).strip() if buffer else clause
        if len(buffer.split()) >= min_words:
            yield buffer
            buffer = ''
    if buffer:
        yield buffer


# ---------- WAV -> PCM frame conversion ----------

def _ffmpeg_resample(wav_bytes: bytes) -> bytes:
    """Convert any WAV payload to raw 48kHz stereo s16le PCM via ffmpeg.

    ffmpeg is already an Echo Core hard dependency (heartbeat, video
    transcription). Using a stdin/stdout pipe keeps the synth tempfile
    short-lived and avoids extra disk traffic.
    """
    proc = subprocess.run(
        [
            'ffmpeg', '-hide_banner', '-loglevel', 'error',
            '-i', 'pipe:0',
            '-f', 's16le',
            '-ar', str(TARGET_SAMPLE_RATE),
            '-ac', str(TARGET_CHANNELS),
            '-acodec', 'pcm_s16le',
            'pipe:1',
        ],
        input=wav_bytes,
        capture_output=True,
        check=False,
    )
    if proc.returncode != 0:
        err = proc.stderr.decode('utf-8', errors='replace')[:200]
        raise RuntimeError(f"ffmpeg resample failed (rc={proc.returncode}): {err}")
    return proc.stdout


def _is_target_format(wav_bytes: bytes) -> bool:
    """Quick check whether the WAV already matches Discord's PCM format."""
    try:
        with wave.open(io.BytesIO(wav_bytes), 'rb') as w:
            return (
                w.getframerate() == TARGET_SAMPLE_RATE
                and w.getnchannels() == TARGET_CHANNELS
                and w.getsampwidth() == TARGET_SAMPLE_WIDTH
                and w.getcomptype() == 'NONE'
            )
    except (wave.Error, EOFError):
        return False


def _extract_pcm_native(wav_bytes: bytes) -> bytes:
    """Strip the WAV header and return raw PCM (target format assumed)."""
    with wave.open(io.BytesIO(wav_bytes), 'rb') as w:
        return w.readframes(w.getnframes())


def wav_to_pcm_20ms_frames(wav_bytes: bytes) -> List[bytes]:
    """Parse a WAV blob, normalize to 48kHz s16le stereo, slice into 20ms frames.

    The final frame is zero-padded to a full 3840 bytes so Discord's audio
    thread always reads whole frames. Empty input yields an empty list.
    """
    if not wav_bytes:
        return []
    pcm = _extract_pcm_native(wav_bytes) if _is_target_format(wav_bytes) else _ffmpeg_resample(wav_bytes)
    if not pcm:
        return []
    frames: List[bytes] = []
    for offset in range(0, len(pcm), FRAME_BYTES):
        chunk = pcm[offset:offset + FRAME_BYTES]
        if len(chunk) < FRAME_BYTES:
            chunk = chunk + b'\x00' * (FRAME_BYTES - len(chunk))
        frames.append(chunk)
    return frames


# ---------- TTS worker queue ----------

class TTSQueue:
    """Worker thread: text in -> 20ms PCM frames out.

    Usage::

        ttsq = TTSQueue(voice_id="M2", lang="ro")
        ttsq.start()
        ttsq.push_text("salut Marius, ce mai faci?")
        voice_client.play(EchoStreamingAudioSource(ttsq))
        # ... barge-in detected:
        ttsq.clear()
        # ... session over:
        ttsq.stop()
    """

    def __init__(self, voice_id: str = "M2", lang: str = "ro"):
        self.voice_id = voice_id
        self.lang = lang
        self._text_queue: queue.Queue = queue.Queue()
        self._pcm_queue: queue.Queue = queue.Queue()
        self._worker_thread: Optional[threading.Thread] = None
        self._stop_event = threading.Event()

    # --- lifecycle ---

    def start(self) -> None:
        if self._worker_thread is not None and self._worker_thread.is_alive():
            return
        self._stop_event.clear()
        self._worker_thread = threading.Thread(
            target=self._worker_loop,
            name=f"tts-worker-{self.voice_id}",
            daemon=True,
        )
        self._worker_thread.start()

    def stop(self) -> None:
        """Signal the worker to exit, drain queues, join (timeout 5s)."""
        self._stop_event.set()
        # Wake the worker if it's blocked on get(timeout=...).
        self._text_queue.put(_POISON)
        thread = self._worker_thread
        if thread is not None:
            thread.join(timeout=5.0)
        self._worker_thread = None
        self._drain(self._text_queue)
        self._drain(self._pcm_queue)

    # --- producer side ---

    def push_text(self, text: str) -> None:
        """Normalize, segment into clauses, enqueue each clause for synthesis."""
        if not text:
            return
        cleaned = normalize_for_tts(text)
        n = 0
        for clause in clause_segments(cleaned):
            clause = clause.strip()
            if clause:
                self._text_queue.put(clause)
                n += 1
        log.info("ttsq.push_text: input %d chars → %d clauses queued", len(text), n)

    def clear(self) -> None:
        """Drop everything pending (used for barge-in)."""
        self._drain(self._text_queue)
        self._drain(self._pcm_queue)

    def is_empty(self) -> bool:
        return self._text_queue.empty() and self._pcm_queue.empty()

    # --- consumer side (called by EchoStreamingAudioSource) ---

    def get_frame_nowait(self) -> Optional[bytes]:
        """Return the next PCM frame if available, else None — no blocking.

        Blocking inside the player's read() loop wrecks Discord's 20ms cadence
        and the client interprets the stream as stuttering / out-of-order.
        """
        try:
            return self._pcm_queue.get_nowait()
        except queue.Empty:
            return None

    # --- internals ---

    @staticmethod
    def _drain(q: queue.Queue) -> None:
        while True:
            try:
                q.get_nowait()
            except queue.Empty:
                return

    def _worker_loop(self) -> None:
        while not self._stop_event.is_set():
            try:
                item = self._text_queue.get(timeout=0.1)
            except queue.Empty:
                continue
            if item is _POISON:
                break
            if not isinstance(item, str):
                continue
            preview = item[:60]
            try:
                result = synthesize(item, voice=self.voice_id, lang=self.lang)
            except Exception as e:
                log.warning("TTS synth raised for %r: %s", preview, e)
                continue
            if not result.get('ok'):
                log.warning("TTS synth not ok for %r: %s", preview, result.get('error'))
                continue
            path = result.get('path')
            if not path:
                log.warning("TTS synth ok but no path for %r", preview)
                continue
            wav_bytes = b''
            try:
                wav_bytes = Path(path).read_bytes()
            except OSError as e:
                log.warning("TTS WAV read failed for %r: %s", preview, e)
            finally:
                try:
                    Path(path).unlink(missing_ok=True)
                except OSError:
                    pass
            if not wav_bytes:
                continue
            try:
                frames = wav_to_pcm_20ms_frames(wav_bytes)
            except RuntimeError as e:
                log.warning("TTS WAV-to-PCM failed for %r: %s", preview, e)
                continue
            if not frames:
                log.warning("TTS WAV-to-PCM produced 0 frames for %r", preview)
                continue
            for frame in frames:
                if self._stop_event.is_set():
                    return
                self._pcm_queue.put(frame)
            log.info("TTS pushed %d frames (%.1fs) for %r",
                     len(frames), len(frames) * 0.02, preview)


# ---------- Discord audio source ----------

class EchoStreamingAudioSource(discord.AudioSource):
    """Pull PCM frames from a ``TTSQueue`` into Discord's audio thread.

    A single ``voice_client.play(EchoStreamingAudioSource(ttsq))`` call
    spans the whole session. When the TTS queue is empty, ``read()``
    returns a 20ms silence frame to keep the player alive — otherwise
    Discord would interpret an empty return as end-of-stream and stop
    the player, so real TTS frames pushed later would be silently
    discarded. The player is explicitly terminated only via
    ``cleanup()`` (called on voice session teardown).
    """

    # 20ms of s16le stereo at 48kHz silence (960 samples × 2 channels × 2 bytes).
    _SILENCE_FRAME = b'\x00' * (960 * 2 * 2)

    def __init__(self, ttsq: TTSQueue):
        self._ttsq = ttsq
        self._closed = False

    def read(self) -> bytes:
        if self._closed:
            return b''
        frame = self._ttsq.get_frame_nowait()
        if frame is None:
            return self._SILENCE_FRAME
        return frame

    def is_opus(self) -> bool:
        return False

    def cleanup(self) -> None:
        self._closed = True
        try:
            self._ttsq.clear()
        except Exception:
            pass
