"""WebSocket streaming endpoint with server-side VAD and barge-in support."""

from __future__ import annotations

import asyncio
import base64
import json
import logging
import time
import uuid

import numpy as np
from fastapi import WebSocket, WebSocketDisconnect

from roohai.audio import audio_to_wav_bytes
from roohai.pipeline import Rooh

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__)


async def handle_streaming(websocket: WebSocket, pipeline: Rooh) -> None:
    await websocket.accept()
    session_id = str(uuid.uuid4())
    logger.info("WebSocket session started: %s", session_id)
    await websocket.send_json({"type": "session", "session_id": session_id})

    if pipeline.stt_supports_streaming:
        await _handle_streaming_stt(websocket, pipeline, session_id)
    else:
        await _handle_batch_stt(websocket, pipeline, session_id)


def _ws_send_interrupt(websocket: WebSocket):
    """Return a send_interrupt callback bound to the given WebSocket."""
    async def _send():
        await websocket.send_json({"type": "interrupt"})
    return _send


async def _ws_send_audio(websocket: WebSocket, wav_bytes: bytes) -> None:
    """Send WAV bytes over WebSocket as base64-encoded JSON."""
    audio_b64 = base64.b64encode(wav_bytes).decode()
    await websocket.send_json({"type": "audio", "audio": audio_b64})


async def _handle_streaming_stt(websocket: WebSocket, pipeline: Rooh, session_id: str) -> None:
    """Streaming STT path: forward audio to Deepgram, use local VAD for UI only."""
    vad = pipeline.vad
    stt = pipeline.stt
    is_speaking = False
    paused = False
    state = BargeInState()
    pending_tasks: set[asyncio.Task] = set()

    async def on_transcript(text: str, is_final: bool) -> None:
        """Called by STT with interim/final segment transcripts."""
        try:
            await websocket.send_json({
                "type": "interim_transcription",
                "text": text,
            })
        except Exception:
            logger.debug("Failed to send transcript to client", exc_info=True)

    async def _run_generation(text: str) -> None:
        """Run a single LLM+TTS generation with barge-in support."""
        gen_id = await state.start_generation(_ws_send_interrupt(websocket))
        try:
            await _process_transcript(websocket, pipeline, text, state.cancel_event, session_id=session_id)
        finally:
            state.finish_generation()

    async def on_utterance_end(text: str) -> None:
        """Called when STT detects end of turn -- fire off LLM+TTS as a task."""
        task = asyncio.create_task(_run_generation(text))
        pending_tasks.add(task)
        task.add_done_callback(pending_tasks.discard)

    try:
        # Start streaming STT connection
        await stt.start_stream(on_transcript, on_utterance_end)
        await websocket.send_json({"type": "status", "state": "listening"})

        while True:
            msg = await websocket.receive_text()
            data = json.loads(msg)

            if data.get("type") == "audio_chunk":
                if paused:
                    continue

                raw = base64.b64decode(data["audio"])

                # Forward raw PCM16 bytes directly to STT
                await stt.send_chunk(raw)

                # Run local VAD for UI indicator + barge-in detection
                if vad is not None:
                    chunk = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
                    speech_prob = await asyncio.to_thread(vad.is_speech, chunk, 16000)

                    if speech_prob >= SPEECH_THRESHOLD:
                        if not is_speaking:
                            is_speaking = True
                            await websocket.send_json(
                                {"type": "status", "state": "listening", "vad": "speech_start"}
                            )
                        # Barge-in: check on every speech chunk so we catch it
                        # even if the pipeline finishes between chunks.
                        if state.is_active and not state.cancel_event.is_set():
                            logger.info(
                                "Barge-in detected via VAD during generation %d",
                                state.generation_id,
                            )
                            state.cancel_event.set()
                            try:
                                await websocket.send_json({"type": "interrupt"})
                            except Exception:
                                pass
                    elif speech_prob < SPEECH_THRESHOLD and is_speaking:
                        is_speaking = False
                        await websocket.send_json(
                            {"type": "status", "state": "listening", "vad": "speech_end"}
                        )

            elif data.get("type") == "pause_audio":
                paused = True
                logger.info("Audio paused by client")
                await websocket.send_json({"type": "status", "state": "paused"})

            elif data.get("type") == "resume_audio":
                paused = False
                logger.info("Audio resumed by client")
                await websocket.send_json({"type": "status", "state": "listening"})

            elif data.get("type") == "playback_done":
                state.playback_done()

            elif data.get("type") == "ping":
                await websocket.send_json({"type": "pong"})

    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.exception("WebSocket streaming error")
        try:
            await websocket.send_json({"type": "error", "message": str(e)})
        except Exception:
            pass
    finally:
        for task in pending_tasks:
            task.cancel()
        await asyncio.gather(*pending_tasks, return_exceptions=True)
        pending_tasks.clear()
        await stt.stop_stream()


async def _handle_batch_stt(websocket: WebSocket, pipeline: Rooh, session_id: str) -> None:
    """Batch STT path: accumulate speech via VAD, then transcribe full utterance."""
    vad = pipeline.vad
    if vad is None:
        await websocket.send_json({"type": "error", "message": "VAD model not loaded"})
        await websocket.close()
        return

    speech_buffer: list[np.ndarray] = []
    is_speaking = False
    paused = False
    silence_count = 0
    speech_chunk_count = 0
    state = BargeInState()
    pending_tasks: set[asyncio.Task] = set()

    try:
        await websocket.send_json({"type": "status", "state": "listening"})

        while True:
            msg = await websocket.receive_text()
            data = json.loads(msg)

            if data.get("type") == "audio_chunk":
                if paused:
                    continue

                # Decode PCM16 audio chunk
                raw = base64.b64decode(data["audio"])
                chunk = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0

                # Run VAD
                speech_prob = await asyncio.to_thread(vad.is_speech, chunk, 16000)

                # Barge-in: speech detected while AI is active.
                # Check on EVERY speech chunk (not just the first) so that
                # barge-in still fires if the pipeline finishes between chunks.
                if speech_prob >= SPEECH_THRESHOLD and state.is_active:
                    if not state.cancel_event.is_set():
                        logger.info("Barge-in detected via VAD during generation %d", state.generation_id)
                        state.cancel_event.set()
                        state.playback_pending = False
                        try:
                            await websocket.send_json({"type": "interrupt"})
                        except Exception:
                            pass

                if speech_prob >= SPEECH_THRESHOLD:
                    if not is_speaking:
                        is_speaking = True
                        await websocket.send_json(
                            {"type": "status", "state": "listening", "vad": "speech_start"}
                        )
                    speech_buffer.append(chunk)
                    speech_chunk_count += 1
                    silence_count = 0
                else:
                    if is_speaking:
                        speech_buffer.append(chunk)
                        silence_count += 1

                    if is_speaking and (
                        (silence_count >= SILENCE_CHUNKS_TO_END and speech_chunk_count >= MIN_SPEECH_CHUNKS)
                        or len(speech_buffer) >= MAX_SPEECH_BUFFER_CHUNKS
                    ):
                        # Speech ended or buffer overflow -- process the utterance
                        if len(speech_buffer) >= MAX_SPEECH_BUFFER_CHUNKS:
                            logger.warning(
                                "Speech buffer overflow (%d chunks) — processing early",
                                len(speech_buffer),
                            )
                        await websocket.send_json(
                            {"type": "status", "state": "listening", "vad": "speech_end"}
                        )
                        audio = np.concatenate(speech_buffer)
                        speech_buffer = []
                        is_speaking = False
                        silence_count = 0
                        speech_chunk_count = 0

                        task = asyncio.create_task(
                            _process_utterance(websocket, pipeline, audio, state, session_id=session_id)
                        )
                        pending_tasks.add(task)
                        task.add_done_callback(pending_tasks.discard)

            elif data.get("type") == "pause_audio":
                paused = True
                logger.info("Audio paused by client (batch mode)")
                await websocket.send_json({"type": "status", "state": "paused"})

            elif data.get("type") == "resume_audio":
                paused = False
                logger.info("Audio resumed by client (batch mode)")
                await websocket.send_json({"type": "status", "state": "listening"})

            elif data.get("type") == "playback_done":
                state.playback_done()

            elif data.get("type") == "ping":
                await websocket.send_json({"type": "pong"})

    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.exception("WebSocket streaming error")
        try:
            await websocket.send_json({"type": "error", "message": str(e)})
        except Exception:
            pass
    finally:
        for task in pending_tasks:
            task.cancel()
        await asyncio.gather(*pending_tasks, return_exceptions=True)
        pending_tasks.clear()


async def _process_transcript(
    websocket: WebSocket,
    pipeline: Rooh,
    transcript: str,
    cancel: asyncio.Event | None = None,
    session_id: str | None = None,
) -> None:
    """Run LLM -> TTS on a completed transcript (streaming or batch per config)."""
    try:
        if not transcript or not transcript.strip():
            await websocket.send_json({"type": "status", "state": "listening"})
            return

        metrics = PipelineMetrics()

        # Send the final committed transcription (replaces any interim)
        await websocket.send_json({"type": "transcription", "text": transcript})

        if cancel and cancel.is_set():
            logger.info("Generation cancelled before LLM")
            return

        await websocket.send_json({"type": "status", "state": "thinking"})

        async def send_audio(wav_bytes: bytes) -> None:
            await _ws_send_audio(websocket, wav_bytes)

        async def send_chunk(token: str) -> None:
            await websocket.send_json({"type": "llm_response_chunk", "text": token})

        async def send_done(full_text: str) -> None:
            await websocket.send_json({"type": "llm_response_done", "text": full_text})

        async def send_full(full_text: str) -> None:
            await websocket.send_json({"type": "llm_response", "text": full_text})

        await websocket.send_json({"type": "status", "state": "speaking"})

        await generate_and_synthesize(
            pipeline, transcript, send_audio, send_chunk, send_done, send_full,
            cancel, metrics=metrics, session_id=session_id,
        )

        metrics.finish()
        logger.info("Pipeline metrics: %s", metrics.to_dict())
        if not (cancel and cancel.is_set()):
            await websocket.send_json({"type": "metrics", **metrics.to_dict()})
            await websocket.send_json({"type": "status", "state": "listening"})
    except Exception as e:
        logger.exception("Transcript processing failed")
        try:
            await websocket.send_json({"type": "error", "message": str(e)})
            await websocket.send_json({"type": "status", "state": "listening"})
        except Exception:
            pass


async def _process_utterance(
    websocket: WebSocket,
    pipeline: Rooh,
    audio: np.ndarray,
    state: BargeInState,
    session_id: str | None = None,
) -> None:
    """Run STT -> LLM -> TTS on the detected utterance (streaming or batch per config)."""
    gen_id = await state.start_generation(_ws_send_interrupt(websocket))
    cancel = state.cancel_event
    try:
        metrics = PipelineMetrics()

        # STT
        await websocket.send_json({"type": "status", "state": "transcribing"})
        stt_start = time.monotonic()
        wav_bytes = audio_to_wav_bytes(audio, 16000)
        transcription = await pipeline.transcribe(wav_bytes)
        metrics.stt_ms = _elapsed_ms(stt_start)
        await websocket.send_json({"type": "transcription", "text": transcription})

        if not transcription or not transcription.strip():
            await websocket.send_json({"type": "status", "state": "listening"})
            return

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

        await websocket.send_json({"type": "status", "state": "thinking"})

        async def send_audio(wav_bytes: bytes) -> None:
            await _ws_send_audio(websocket, wav_bytes)

        async def send_chunk(token: str) -> None:
            await websocket.send_json({"type": "llm_response_chunk", "text": token})

        async def send_done(full_text: str) -> None:
            await websocket.send_json({"type": "llm_response_done", "text": full_text})

        async def send_full(full_text: str) -> None:
            await websocket.send_json({"type": "llm_response", "text": full_text})

        await websocket.send_json({"type": "status", "state": "speaking"})

        await generate_and_synthesize(
            pipeline, transcription, send_audio, send_chunk, send_done, send_full,
            cancel, metrics=metrics, session_id=session_id,
        )

        metrics.finish()
        logger.info("Pipeline metrics: %s", metrics.to_dict())
        if not cancel.is_set():
            await websocket.send_json({"type": "metrics", **metrics.to_dict()})
            await websocket.send_json({"type": "status", "state": "listening"})
    except Exception as e:
        logger.exception("Utterance processing failed")
        try:
            await websocket.send_json({"type": "error", "message": str(e)})
            await websocket.send_json({"type": "status", "state": "listening"})
        except Exception:
            pass
    finally:
        state.finish_generation()
