"""Shared voice pipeline logic for WebSocket and WebRTC transports.

Provides barge-in state management, VAD constants, TTS sentence splitting,
and transport-agnostic TTS chunking.
"""

from __future__ import annotations

import asyncio
import logging
import re
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field

from roohai.pipeline import Rooh

logger = logging.getLogger(__name__)


@dataclass
class PipelineMetrics:
    """Timing metrics for a single pipeline turn (STT -> LLM -> TTS)."""

    stt_ms: float | None = None
    llm_ms: float | None = None
    llm_first_token_ms: float | None = None  # streaming only
    tts_ms: float | None = None
    total_ms: float | None = None

    # Internal timestamps (not sent to client)
    _start: float = field(default_factory=time.monotonic, repr=False)

    def finish(self) -> None:
        self.total_ms = _elapsed_ms(self._start)

    def to_dict(self) -> dict:
        d: dict = {}
        if self.stt_ms is not None:
            d["stt_ms"] = round(self.stt_ms, 1)
        if self.llm_ms is not None:
            d["llm_ms"] = round(self.llm_ms, 1)
        if self.llm_first_token_ms is not None:
            d["llm_first_token_ms"] = round(self.llm_first_token_ms, 1)
        if self.tts_ms is not None:
            d["tts_ms"] = round(self.tts_ms, 1)
        if self.total_ms is not None:
            d["total_ms"] = round(self.total_ms, 1)
        return d


def _elapsed_ms(start: float) -> float:
    return (time.monotonic() - start) * 1000

# VAD parameters
CHUNK_SIZE = 512  # samples per chunk at 16kHz (~32ms)
SPEECH_THRESHOLD = 0.5
SILENCE_CHUNKS_TO_END = 30  # ~1 second of silence to end utterance
MIN_SPEECH_CHUNKS = 10  # minimum ~320ms of speech
MAX_SPEECH_BUFFER_CHUNKS = 600  # ~19s at 512 samples/32ms — cap to prevent unbounded growth

# Maximum text length for a single TTS call (chars). Longer text is split into sentences.
# SpeechT5 can handle ~100-150 words (~600 tokens). Keep chunks short for reliability.
TTS_MAX_CHARS = 250


class BargeInState:
    """Per-connection state for barge-in / interrupt support.

    Tracks two phases:
      1. ``is_generating`` -- server is running LLM + TTS (can cancel work)
      2. ``playback_pending`` -- client is still playing queued audio

    Barge-in triggers during *either* phase.
    """

    def __init__(self) -> None:
        self.generation_id: int = 0
        self.cancel_event: asyncio.Event = asyncio.Event()
        self.is_generating: bool = False
        self.playback_pending: bool = False
        self._gen_done: asyncio.Event = asyncio.Event()
        self._gen_done.set()  # no generation in progress initially

    @property
    def is_active(self) -> bool:
        """True while the AI is generating OR the client is still playing audio."""
        return self.is_generating or self.playback_pending

    async def start_generation(self, send_interrupt: Callable[[], Awaitable[None]]) -> int:
        """Begin a new generation, cancelling any in-flight one first.

        Args:
            send_interrupt: Async callback that sends an interrupt signal to the client.
        """
        if self.is_active:
            logger.info("Barge-in detected — cancelling generation %d", self.generation_id)
            self.cancel_event.set()
            self.playback_pending = False
            try:
                await send_interrupt()
            except Exception:
                pass
            # Wait for the in-flight generation to exit (with timeout)
            if self.is_generating:
                try:
                    await asyncio.wait_for(self._gen_done.wait(), timeout=2.0)
                except asyncio.TimeoutError:
                    logger.warning("Timed out waiting for generation %d to cancel", self.generation_id)

        self.generation_id += 1
        self.cancel_event = asyncio.Event()
        self._gen_done = asyncio.Event()
        self.is_generating = True
        self.playback_pending = True
        return self.generation_id

    def finish_generation(self) -> None:
        """Mark server-side generation as done (client may still be playing)."""
        self.is_generating = False
        self._gen_done.set()
        # playback_pending stays True -- cleared by client's playback_done message

    def playback_done(self) -> None:
        """Called when the client reports it finished playing all audio."""
        self.playback_pending = False


def split_into_sentences(text: str) -> list[str]:
    """Split text into sentences for TTS chunking."""
    # Split on sentence-ending punctuation followed by space or end-of-string
    parts = re.split(r'(?<=[.!?])\s+', text.strip())
    # Merge very short fragments with previous
    sentences: list[str] = []
    for part in parts:
        if sentences and len(sentences[-1]) < 40:
            sentences[-1] += " " + part
        else:
            sentences.append(part)
    return [s for s in sentences if s.strip()]


SENTENCE_BOUNDARY_RE = re.compile(r'(?<=[.!?])\s')


async def generate_and_synthesize(
    pipeline: Rooh,
    text: str,
    send_audio: Callable[[bytes], Awaitable[None]],
    send_chunk: Callable[[str], Awaitable[None]],
    send_done: Callable[[str], Awaitable[None]],
    send_full: Callable[[str], Awaitable[None]],
    cancel: asyncio.Event | None = None,
    history: list[dict] | None = None,
    system_prompt: str | None = None,
    metrics: PipelineMetrics | None = None,
    session_id: str | None = None,
) -> str:
    """Run LLM + TTS, choosing streaming or batch based on ``pipeline.llm_streaming``.

    Args:
        send_chunk: Called per token when streaming.
        send_done: Called with full text when streaming finishes.
        send_full: Called with full text when using batch (non-streaming) mode.
        metrics: Optional metrics object to populate with LLM/TTS timing.
        session_id: Unique connection identifier passed through to the LLM/hook.
    """
    if pipeline.llm_streaming:
        return await stream_and_synthesize(
            pipeline, text, send_audio, send_chunk, send_done,
            cancel, history, system_prompt, metrics, session_id,
        )
    else:
        return await _batch_and_synthesize(
            pipeline, text, send_audio, send_full, cancel, history, system_prompt, metrics, session_id,
        )


async def _batch_and_synthesize(
    pipeline: Rooh,
    text: str,
    send_audio: Callable[[bytes], Awaitable[None]],
    send_full: Callable[[str], Awaitable[None]],
    cancel: asyncio.Event | None = None,
    history: list[dict] | None = None,
    system_prompt: str | None = None,
    metrics: PipelineMetrics | None = None,
    session_id: str | None = None,
) -> str:
    """Non-streaming LLM: wait for full response, then synthesize."""
    llm_start = time.monotonic()
    response = await pipeline.chat(text, history=history, system_prompt=system_prompt, session_id=session_id)
    if metrics:
        metrics.llm_ms = _elapsed_ms(llm_start)
    if cancel and cancel.is_set():
        return response
    await send_full(response)
    tts_start = time.monotonic()
    await synthesize_chunks(pipeline, response, send_audio, cancel)
    if metrics:
        metrics.tts_ms = _elapsed_ms(tts_start)
    return response


async def stream_and_synthesize(
    pipeline: Rooh,
    text: str,
    send_audio: Callable[[bytes], Awaitable[None]],
    send_chunk: Callable[[str], Awaitable[None]],
    send_done: Callable[[str], Awaitable[None]],
    cancel: asyncio.Event | None = None,
    history: list[dict] | None = None,
    system_prompt: str | None = None,
    metrics: PipelineMetrics | None = None,
    session_id: str | None = None,
) -> str:
    """Stream LLM response while synthesizing completed sentences.

    - Consumes ``pipeline.chat_stream()``
    - Sends incremental ``llm_response_chunk`` via *send_chunk*
    - When a sentence boundary is detected, starts TTS on that sentence
    - Returns the full accumulated response text
    """
    full_text = ""
    sentence_buf = ""
    llm_start = time.monotonic()
    first_token_recorded = False
    tts_total = 0.0

    async def _synthesize_sentence(sentence: str) -> None:
        nonlocal tts_total
        if cancel and cancel.is_set():
            return
        tts_start = time.monotonic()
        await synthesize_chunks(pipeline, sentence, send_audio, cancel)
        tts_total += _elapsed_ms(tts_start)

    async for token in pipeline.chat_stream(text, history=history, system_prompt=system_prompt, session_id=session_id):
        if cancel and cancel.is_set():
            break
        if not first_token_recorded and metrics:
            metrics.llm_first_token_ms = _elapsed_ms(llm_start)
            first_token_recorded = True
        full_text += token
        sentence_buf += token
        await send_chunk(token)

        # Check for sentence boundary in the buffer
        match = SENTENCE_BOUNDARY_RE.search(sentence_buf)
        if match:
            # Extract completed sentence(s) up to the boundary
            split_pos = match.end()
            completed = sentence_buf[:split_pos].strip()
            sentence_buf = sentence_buf[split_pos:]
            if completed:
                await _synthesize_sentence(completed)

    if metrics:
        metrics.llm_ms = _elapsed_ms(llm_start) - tts_total  # exclude TTS time interleaved

    # Synthesize any remaining text
    remainder = sentence_buf.strip()
    if remainder and not (cancel and cancel.is_set()):
        await _synthesize_sentence(remainder)

    if metrics:
        metrics.tts_ms = tts_total

    await send_done(full_text)
    return full_text


async def synthesize_chunks(
    pipeline: Rooh,
    text: str,
    send_audio: Callable[[bytes], Awaitable[None]],
    cancel: asyncio.Event | None = None,
) -> None:
    """Synthesize text to speech, splitting into chunks if needed.

    Transport-agnostic: uses *send_audio* callback to deliver WAV bytes.
    If *cancel* is set at any point, remaining chunks are skipped.
    """
    if cancel and cancel.is_set():
        return

    if len(text) <= TTS_MAX_CHARS:
        tts_wav = await pipeline.synthesize(text)
        if cancel and cancel.is_set():
            return
        await send_audio(tts_wav)
    else:
        sentences = split_into_sentences(text)
        batch = ""
        for sentence in sentences:
            if cancel and cancel.is_set():
                return
            if batch and len(batch) + len(sentence) + 1 > TTS_MAX_CHARS:
                tts_wav = await pipeline.synthesize(batch)
                if cancel and cancel.is_set():
                    return
                await send_audio(tts_wav)
                batch = sentence
            else:
                batch = (batch + " " + sentence).strip() if batch else sentence
        if batch:
            if cancel and cancel.is_set():
                return
            tts_wav = await pipeline.synthesize(batch)
            if cancel and cancel.is_set():
                return
            await send_audio(tts_wav)
