"""Custom audio tracks for WebRTC processing with full barge-in and VAD support."""

import asyncio
import io
import json
import logging
import time
from fractions import Fraction

import numpy as np
from aiortc import MediaStreamTrack
from av import AudioFrame

from roohai.audio import audio_to_wav_bytes

from .voice_common import (
    CHUNK_SIZE,
    MAX_SPEECH_BUFFER_CHUNKS,
    MIN_SPEECH_CHUNKS,
    SILENCE_CHUNKS_TO_END,
    SPEECH_THRESHOLD,
    BargeInState,
    PipelineMetrics,
    _elapsed_ms,
    generate_and_synthesize,
    synthesize_chunks,
)

logger = logging.getLogger(__name__)

# WebRTC audio parameters
SAMPLE_RATE = 48000  # WebRTC default
FRAME_SAMPLES = 960  # 20ms at 48kHz
AUDIO_PTIME = FRAME_SAMPLES / SAMPLE_RATE  # 0.02s = 20ms
TARGET_SR = 16000  # Pipeline expects 16kHz
PRE_SPEECH_CHUNKS = 6  # ~192ms of audio kept before VAD triggers


class AudioProcessorTrack(MediaStreamTrack):
    """Receives audio from client, processes through pipeline, returns TTS audio."""

    kind = "audio"

    def __init__(self, pipeline, data_channel=None, session_id: str | None = None):
        super().__init__()
        self.track = None  # set later via set_input_track()
        self.pipeline = pipeline
        self.data_channel = data_channel
        self.session_id = session_id
        self._buffer: list[np.ndarray] = []  # 16kHz float32 chunks for batch STT
        self._response_frames: asyncio.Queue[AudioFrame] = asyncio.Queue()
        self._barge_in = BargeInState()
        self._running = True
        self._process_task: asyncio.Task | None = None
        self._pipeline_tasks: set[asyncio.Task] = set()
        self._pts = 0

        # recv() pacing — aiortc calls recv() in a tight loop, so we must
        # sleep to emit frames at real-time 20ms intervals (like the base
        # AudioStreamTrack does).  Without this the event loop starves.
        self._recv_start: float | None = None
        self._recv_timestamp = 0

        # Silero VAD state
        self._is_speaking = False
        self._silence_count = 0
        self._speech_chunk_count = 0
        self._vad_accum: list[np.ndarray] = []  # accumulate 16kHz samples to CHUNK_SIZE

        # Pre-speech ring buffer: keep last ~200ms of audio so we don't clip
        # the start of speech.  At 16kHz with 512-sample VAD chunks, that's
        # about 6 chunks (~192ms).
        self._pre_speech_chunks: list[np.ndarray] = []

        # Streaming STT readiness flag — only True after start_stream() succeeds
        self._stt_streaming_ready = False

        # Pause flag — when True, incoming audio is ignored (mic muted by client)
        self._paused = False

    def set_input_track(self, track: MediaStreamTrack) -> None:
        """Attach the remote audio track and start the processing loop."""
        self.track = track
        self._process_task = asyncio.ensure_future(self._process_loop())

    def _send_status(self, msg_type: str, **kwargs):
        """Send a JSON status update to the frontend via data channel."""
        if self.data_channel and self.data_channel.readyState == "open":
            payload = {"type": msg_type, **kwargs}
            self.data_channel.send(json.dumps(payload))

    async def _send_interrupt(self) -> None:
        """Send interrupt via data channel (used as BargeInState callback)."""
        self._send_status("interrupt")
        self._clear_response_queue()

    def _clear_response_queue(self):
        """Drain stale TTS frames from the response queue on barge-in."""
        while not self._response_frames.empty():
            try:
                self._response_frames.get_nowait()
            except asyncio.QueueEmpty:
                break

    @staticmethod
    def _downsample_48k_to_16k(samples_48k: np.ndarray) -> np.ndarray:
        """Downsample 48kHz mono float32 to 16kHz by decimation (take every 3rd sample).

        This matches the browser AudioWorklet approach used in the WebSocket path.
        Speech energy is well below the 8kHz Nyquist limit so simple decimation
        preserves quality without needing an explicit anti-aliasing filter.
        """
        return samples_48k[::3].copy()

    async def _process_loop(self):
        """Continuously receive frames, detect speech via Silero VAD, run pipeline."""
        frame_count = 0
        logger.info("_process_loop started — waiting for audio frames")
        while self._running:
            try:
                frame = await self.track.recv()
            except Exception:
                logger.info("Input track ended after %d frames", frame_count)
                break

            frame_count += 1
            if frame_count == 1:
                logger.info("First audio frame received from WebRTC track")
            elif frame_count % 500 == 0:
                logger.info("Processed %d audio frames so far", frame_count)

            # Convert frame to numpy float32 mono
            raw = frame.to_ndarray()

            if frame_count == 1:
                logger.info(
                    "Frame format=%s layout=%s shape=%s",
                    frame.format.name, frame.layout.name, raw.shape,
                )

            # Handle packed s16 (interleaved stereo) vs planar s16p
            if frame.format.name == "s16" and frame.layout.name == "stereo":
                # Packed interleaved: shape (1, samples*2) → L R L R ...
                data = raw.flatten().astype(np.float64)
                left = data[0::2]
                right = data[1::2]
                samples = (left + right) / 2.0  # proper stereo→mono mix
            elif raw.ndim == 2 and raw.shape[0] > 1:
                # Planar multi-channel: shape (channels, samples)
                samples = raw.mean(axis=0).astype(np.float64)
            elif raw.ndim == 2:
                # Planar mono or packed mono: shape (1, samples)
                samples = raw.flatten().astype(np.float64)
            else:
                samples = raw.astype(np.float64)

            samples_48k = (samples / 32768.0).astype(np.float32)

            # Downsample to 16kHz for VAD and pipeline
            samples_16k = self._downsample_48k_to_16k(samples_48k)

            # If streaming STT is connected and ready, forward chunks
            if self._stt_streaming_ready:
                pcm16 = (np.clip(samples_16k, -1.0, 1.0) * 32767).astype(np.int16)
                await self.pipeline.stt.send_chunk(pcm16.tobytes())

            # Accumulate into CHUNK_SIZE (512 samples) for Silero VAD
            self._vad_accum.append(samples_16k)
            accum = np.concatenate(self._vad_accum)

            while len(accum) >= CHUNK_SIZE:
                vad_chunk = accum[:CHUNK_SIZE]
                accum = accum[CHUNK_SIZE:]
                await self._process_vad_chunk(vad_chunk, samples_16k)

            # Store remainder
            if len(accum) > 0:
                self._vad_accum = [accum]
            else:
                self._vad_accum = []

    def handle_pause_resume(self, msg_type: str) -> None:
        """Handle pause/resume messages from data channel."""
        if msg_type == "pause_audio":
            self._paused = True
            logger.info("Audio paused by client (WebRTC)")
            self._send_status("status", state="paused")
        elif msg_type == "resume_audio":
            self._paused = False
            logger.info("Audio resumed by client (WebRTC)")
            self._send_status("status", state="listening")

    async def _process_vad_chunk(self, vad_chunk: np.ndarray, frame_16k: np.ndarray):
        """Process a single CHUNK_SIZE chunk through Silero VAD."""
        if self._paused:
            return

        vad = self.pipeline.vad
        if vad is None:
            logger.warning("VAD is None, skipping chunk")
            return

        speech_prob = await asyncio.to_thread(vad.is_speech, vad_chunk, TARGET_SR)

        if speech_prob >= SPEECH_THRESHOLD:
            # Barge-in: speech detected while AI is active
            if self._barge_in.is_active and not self._is_speaking:
                if not self._barge_in.cancel_event.is_set():
                    logger.info(
                        "Barge-in detected via VAD during generation %d",
                        self._barge_in.generation_id,
                    )
                    self._barge_in.cancel_event.set()
                    self._send_status("interrupt")
                    self._clear_response_queue()

            if not self._is_speaking:
                self._is_speaking = True
                logger.info("VAD: speech start (prob=%.3f)", speech_prob)
                self._send_status("status", state="listening", vad="speech_start")
                # Prepend pre-speech buffer so we don't clip the start of speech
                if self._pre_speech_chunks:
                    logger.info("Prepending %d pre-speech chunks", len(self._pre_speech_chunks))
                    self._buffer.extend(self._pre_speech_chunks)
                    self._pre_speech_chunks = []

            self._buffer.append(vad_chunk)
            self._speech_chunk_count += 1
            self._silence_count = 0
        else:
            if self._is_speaking:
                self._buffer.append(vad_chunk)
                self._silence_count += 1

        # Check if we should process: normal end-of-speech OR buffer overflow
        if self._is_speaking and (
            (self._silence_count >= SILENCE_CHUNKS_TO_END and self._speech_chunk_count >= MIN_SPEECH_CHUNKS)
            or len(self._buffer) >= MAX_SPEECH_BUFFER_CHUNKS
        ):
            # Speech ended or buffer overflow -- process the utterance
            if len(self._buffer) >= MAX_SPEECH_BUFFER_CHUNKS:
                logger.warning(
                    "Speech buffer overflow (%d chunks) — processing early",
                    len(self._buffer),
                )
            total_samples = sum(len(c) for c in self._buffer)
            duration_ms = int(total_samples / TARGET_SR * 1000)
            logger.info(
                "VAD: speech end — %d chunks, %d ms of audio, sending to STT",
                len(self._buffer), duration_ms,
            )
            self._send_status("status", state="listening", vad="speech_end")
            audio = np.concatenate(self._buffer)
            self._buffer = []
            self._is_speaking = False
            self._silence_count = 0
            self._speech_chunk_count = 0

            if self._stt_streaming_ready:
                # Streaming STT handles utterance detection itself
                pass
            else:
                task = asyncio.ensure_future(self._run_pipeline(audio))
                self._pipeline_tasks.add(task)
                task.add_done_callback(self._pipeline_tasks.discard)
        elif not self._is_speaking:
            # Not speaking — maintain rolling pre-speech buffer
            self._pre_speech_chunks.append(vad_chunk)
            if len(self._pre_speech_chunks) > PRE_SPEECH_CHUNKS:
                self._pre_speech_chunks.pop(0)

    async def _run_pipeline(self, audio_16k: np.ndarray):
        """Run the full STT -> LLM -> TTS pipeline on an utterance."""
        gen_id = await self._barge_in.start_generation(self._send_interrupt)
        cancel = self._barge_in.cancel_event

        self._send_status("status", state="transcribing")
        try:
            metrics = PipelineMetrics()

            # Convert 16kHz float32 mono -> WAV bytes for the pipeline
            wav_bytes = audio_to_wav_bytes(audio_16k, TARGET_SR)

            # STT
            stt_start = time.monotonic()
            transcription = await self.pipeline.transcribe(wav_bytes)
            metrics.stt_ms = _elapsed_ms(stt_start)
            logger.info("Transcription: %s", transcription)
            self._send_status("transcription", text=transcription)

            if not transcription or not transcription.strip():
                self._send_status("status", state="listening")
                return

            if cancel.is_set():
                logger.info("Generation %d cancelled before LLM", gen_id)
                return

            # LLM + TTS (streaming or batch per config)
            self._send_status("status", state="thinking")

            async def send_tts_audio(tts_wav_bytes: bytes) -> None:
                if cancel.is_set():
                    return
                tts_audio, tts_sr = self._wav_bytes_to_audio(tts_wav_bytes)
                await self._enqueue_tts_frames(tts_audio, tts_sr)

            async def send_chunk(token: str) -> None:
                self._send_status("llm_response_chunk", text=token)

            async def send_done(full_text: str) -> None:
                self._send_status("llm_response_done", text=full_text)

            async def send_full(full_text: str) -> None:
                self._send_status("llm_response", text=full_text)

            self._send_status("status", state="speaking")

            await generate_and_synthesize(
                self.pipeline, transcription, send_tts_audio,
                send_chunk, send_done, send_full, cancel,
                metrics=metrics, session_id=self.session_id,
            )

            metrics.finish()
            logger.info("Pipeline metrics: %s", metrics.to_dict())
            if not cancel.is_set():
                self._send_status("metrics", **metrics.to_dict())
                self._send_status("playback_complete")
                self._send_status("status", state="listening")
        except Exception:
            logger.exception("Pipeline error")
            self._send_status("error", message="Pipeline processing failed")
        finally:
            self._barge_in.finish_generation()

    async def _run_pipeline_from_text(self, text: str):
        """Run LLM -> TTS on text from streaming STT (skips STT step)."""
        logger.info("Streaming STT transcription: '%s'", text)
        gen_id = await self._barge_in.start_generation(self._send_interrupt)
        cancel = self._barge_in.cancel_event

        try:
            metrics = PipelineMetrics()
            self._send_status("transcription", text=text)

            if not text or not text.strip():
                logger.info("Empty transcription, returning to listening")
                self._send_status("status", state="listening")
                return

            if cancel.is_set():
                return

            self._send_status("status", state="thinking")

            async def send_tts_audio(tts_wav_bytes: bytes) -> None:
                if cancel.is_set():
                    return
                tts_audio, tts_sr = self._wav_bytes_to_audio(tts_wav_bytes)
                await self._enqueue_tts_frames(tts_audio, tts_sr)

            async def send_chunk(token: str) -> None:
                self._send_status("llm_response_chunk", text=token)

            async def send_done(full_text: str) -> None:
                self._send_status("llm_response_done", text=full_text)

            async def send_full(full_text: str) -> None:
                self._send_status("llm_response", text=full_text)

            self._send_status("status", state="speaking")

            await generate_and_synthesize(
                self.pipeline, text, send_tts_audio,
                send_chunk, send_done, send_full, cancel,
                metrics=metrics, session_id=self.session_id,
            )

            metrics.finish()
            logger.info("Pipeline metrics: %s", metrics.to_dict())
            if not cancel.is_set():
                self._send_status("metrics", **metrics.to_dict())
                self._send_status("playback_complete")
                self._send_status("status", state="listening")
        except Exception:
            logger.exception("Pipeline error (from text)")
            self._send_status("error", message="Pipeline processing failed")
        finally:
            self._barge_in.finish_generation()

    async def start_streaming_stt(self):
        """Initialize streaming STT if the pipeline supports it."""
        if not self.pipeline.stt_supports_streaming:
            logger.info("STT model does not support streaming, using batch mode")
            return

        async def on_transcript(text: str, is_final: bool) -> None:
            logger.info("Deepgram interim (final=%s): '%s'", is_final, text)
            self._send_status("interim_transcription", text=text)

        async def on_utterance_end(text: str) -> None:
            logger.info("Deepgram utterance_end: '%s'", text)
            task = asyncio.ensure_future(self._run_pipeline_from_text(text))
            self._pipeline_tasks.add(task)
            task.add_done_callback(self._pipeline_tasks.discard)

        try:
            await self.pipeline.stt.start_stream(on_transcript, on_utterance_end)
            # Only set the flag AFTER start_stream succeeds —
            # if it fails, the batch path stays active as fallback
            self._stt_streaming_ready = True
            logger.info("Streaming STT started successfully for WebRTC")
        except Exception:
            logger.exception(
                "Failed to start streaming STT — falling back to batch mode"
            )
            self._stt_streaming_ready = False

    @staticmethod
    def _wav_bytes_to_audio(wav_bytes: bytes) -> tuple[np.ndarray, int]:
        """Decode WAV bytes to (float32 numpy array, sample_rate)."""
        import soundfile as sf

        buf = io.BytesIO(wav_bytes)
        data, sr = sf.read(buf, dtype="float32")
        if data.ndim > 1:
            data = data.mean(axis=1)
        return data, sr

    async def _enqueue_tts_frames(self, tts_audio: np.ndarray, tts_sr: int = 24000):
        """Convert TTS audio (mono float32) into 48kHz stereo s16 WebRTC frames."""
        if len(tts_audio) == 0:
            return

        # Resample tts_sr -> 48kHz via linear interpolation
        target_len = int(len(tts_audio) * SAMPLE_RATE / tts_sr)
        indices = np.linspace(0, len(tts_audio) - 1, target_len)
        audio_48k = np.interp(indices, np.arange(len(tts_audio)), tts_audio)

        # Convert to int16
        audio_s16 = np.clip(audio_48k * 32767, -32768, 32767).astype(np.int16)

        # Chunk into frames of FRAME_SAMPLES
        for i in range(0, len(audio_s16), FRAME_SAMPLES):
            chunk = audio_s16[i : i + FRAME_SAMPLES]
            if len(chunk) < FRAME_SAMPLES:
                chunk = np.pad(chunk, (0, FRAME_SAMPLES - len(chunk)))

            # Create stereo frame — Opus encoder requires packed s16
            # Interleave L/R: shape (1, FRAME_SAMPLES * 2)
            interleaved = np.empty(len(chunk) * 2, dtype=np.int16)
            interleaved[0::2] = chunk  # left
            interleaved[1::2] = chunk  # right
            frame = AudioFrame.from_ndarray(interleaved.reshape(1, -1), format="s16", layout="stereo")
            frame.sample_rate = SAMPLE_RATE
            frame.pts = self._pts
            frame.time_base = Fraction(1, 48000)
            self._pts += FRAME_SAMPLES
            await self._response_frames.put(frame)

    async def recv(self) -> AudioFrame:
        """Return the next response audio frame, or silence if idle.

        Pacing is critical: aiortc calls recv() in a tight loop.  We must
        sleep so that frames are emitted at real-time 20 ms intervals,
        matching the 48 kHz / 960-sample cadence.  Without this sleep the
        event loop starves and nothing else (VAD, STT, etc.) can run.
        """
        # --- pacing (same approach as aiortc's AudioStreamTrack) ---
        if self._recv_start is not None:
            self._recv_timestamp += FRAME_SAMPLES
            wait = self._recv_start + (self._recv_timestamp / SAMPLE_RATE) - time.time()
            if wait > 0:
                await asyncio.sleep(wait)
        else:
            self._recv_start = time.time()
            self._recv_timestamp = 0

        # --- return queued TTS frame or silence ---
        try:
            frame = self._response_frames.get_nowait()
            return frame
        except asyncio.QueueEmpty:
            pass

        # Silence — packed s16 stereo: shape (1, FRAME_SAMPLES * 2)
        silence = np.zeros((1, FRAME_SAMPLES * 2), dtype=np.int16)
        frame = AudioFrame.from_ndarray(silence, format="s16", layout="stereo")
        frame.sample_rate = SAMPLE_RATE
        frame.pts = self._pts
        frame.time_base = Fraction(1, 48000)
        self._pts += FRAME_SAMPLES
        return frame

    async def _shutdown(self) -> None:
        """Gracefully shut down the processing loop and clean up resources."""
        self._running = False

        if self._process_task is not None and not self._process_task.done():
            try:
                await asyncio.wait_for(self._process_task, timeout=2.0)
            except asyncio.TimeoutError:
                logger.warning("Process loop did not exit in time — cancelling")
                self._process_task.cancel()
                try:
                    await self._process_task
                except asyncio.CancelledError:
                    pass
            self._process_task = None

        # Cancel any in-flight LLM/TTS pipeline tasks
        for task in self._pipeline_tasks:
            task.cancel()
        if self._pipeline_tasks:
            await asyncio.gather(*self._pipeline_tasks, return_exceptions=True)
            self._pipeline_tasks.clear()

        if self._stt_streaming_ready:
            try:
                await self.pipeline.stt.stop_stream()
            except Exception:
                logger.debug("Error stopping STT stream during shutdown", exc_info=True)
            self._stt_streaming_ready = False

    def stop(self):
        """Stop the processing loop."""
        self._running = False
        asyncio.ensure_future(self._shutdown())
        super().stop()
