"""Voice Stream Adapter — wraps model generation for real-time TTS.

Wraps the model's generate_stream() to emit complete sentences.
First sentence prioritized — minimal tokens before TTS can start.
Voice-optimized generation params: shorter max_tokens, lower temperature.
Latency tracking: measures time-to-first-sentence and time-to-complete.
Pipes sentences directly to TTS engine — speaks as it thinks.
"""

from __future__ import annotations

import logging
import time
from typing import Any, Callable, Iterator

logger = logging.getLogger(__name__)


class VoiceAdapter:
    """Adapts the LLM's token stream into sentence-by-sentence output for TTS.

    Key features:
    - Sentence boundary detection (., !, ?, newlines)
    - First sentence comes out ASAP (sub-500ms target)
    - Latency tracking (time-to-first-sentence, time-to-complete)
    - Pipes sentences directly to TTS engine
    """

    SENTENCE_ENDINGS = [". ", "! ", "? ", ".\n", "!\n", "?\n", ".\t", "!\t", "?\t"]

    def __init__(self) -> None:
        self._stats = {
            "total_sentences": 0,
            "total_generations": 0,
            "avg_time_to_first_sentence_s": 0.0,
            "avg_total_time_s": 0.0,
        }

    def stream_sentences(
        self,
        token_stream: Iterator[str],
        on_sentence: Callable[[str], None] | None = None,
        on_first_sentence: Callable[[str, float], None] | None = None,
    ) -> Iterator[str]:
        """Convert a token stream into a sentence stream.

        Args:
            token_stream: iterator yielding text chunks (from model.generate_stream)
            on_sentence: called for each complete sentence
            on_first_sentence: called with first sentence and time-to-first
        Yields:
            Complete sentences as they're generated
        """
        self._stats["total_generations"] += 1
        t0 = time.time()
        first_sentence_time = None
        buffer = ""
        sentence_count = 0

        for chunk in token_stream:
            buffer += chunk

            # Extract complete sentences
            while buffer:
                end_idx = self._find_sentence_end(buffer)
                if end_idx > 0:
                    sentence = buffer[:end_idx]
                    buffer = buffer[end_idx:]

                    sentence_count += 1
                    self._stats["total_sentences"] += 1

                    if sentence_count == 1:
                        first_sentence_time = time.time() - t0
                        self._stats["avg_time_to_first_sentence_s"] = (
                            (self._stats["avg_time_to_first_sentence_s"] * (self._stats["total_generations"] - 1) + first_sentence_time)
                            / self._stats["total_generations"]
                        )
                        if on_first_sentence:
                            on_first_sentence(sentence, first_sentence_time)

                    if on_sentence:
                        on_sentence(sentence)

                    yield sentence
                else:
                    break

        # Flush remaining buffer
        if buffer.strip():
            sentence_count += 1
            self._stats["total_sentences"] += 1
            if on_sentence:
                on_sentence(buffer)
            yield buffer

        total_time = time.time() - t0
        self._stats["avg_total_time_s"] = (
            (self._stats["avg_total_time_s"] * (self._stats["total_generations"] - 1) + total_time)
            / self._stats["total_generations"]
        )

    def _find_sentence_end(self, text: str) -> int:
        """Find the end index of the first complete sentence in text."""
        earliest = -1
        for ending in self.SENTENCE_ENDINGS:
            idx = text.find(ending)
            if idx >= 0:
                end = idx + len(ending)
                if earliest < 0 or end < earliest:
                    earliest = end
        return earliest

    def get_stats(self) -> dict[str, Any]:
        return {
            **self._stats,
            "avg_time_to_first_sentence_ms": round(self._stats["avg_time_to_first_sentence_s"] * 1000, 1),
        }