"""voxturn end-of-turn detection for Pipecat: a drop-in replacement for Pipecat's Smart Turn analyzer.

    from pipecat_voxturn import VoxTurnAnalyzer
    from pipecat.turns.user_stop import TurnAnalyzerUserTurnStopStrategy
    from pipecat.turns.user_turn_strategies import UserTurnStrategies

    strategies = UserTurnStrategies(stop=[TurnAnalyzerUserTurnStopStrategy(turn_analyzer=VoxTurnAnalyzer())])

It plugs into Pipecat the way LocalSmartTurnAnalyzerV3 does (BaseSmartTurn: the same audio buffering, the same call
when VAD reports the user stopped, the same stop_secs timeout). One addition: Smart Turn is asked once, when VAD
reports the stop; if it says "incomplete" the user waits for the stop_secs timeout (3 s by default). voxturn becomes
more confident as the silence goes on, so while the user stays silent it is asked again every `recheck_secs` and the
turn ends as soon as P(end of turn) reaches `threshold`. Set recheck_secs=None for Smart Turn's single check.
Re-checks run on the event loop (~8 ms on a recent GPU, every 250 ms by default).
"""
import threading
from typing import Any

import numpy as np
from loguru import logger

from pipecat.audio.turn.base_turn_analyzer import EndOfTurnState
from pipecat.audio.turn.smart_turn.base_smart_turn import BaseSmartTurn, SmartTurnParams
from pipecat.metrics.metrics import MetricsData

DEFAULT_MODEL = "AustinXiao/voxturn-v1"


class VoxTurnParams(SmartTurnParams):
    """SmartTurnParams plus voxturn's settings.

    Parameters:
        threshold: P(end of turn) at which the turn is complete.
        recheck_secs: While the user stays silent after VAD reports the stop, ask again this often (seconds);
            None asks only once, like Smart Turn.
        lang: Prompt language, "en" or "zh". Any spoken language works with either.
    """

    threshold: float = 0.8
    recheck_secs: float | None = 0.25
    lang: str = "en"


class VoxTurnAnalyzer(BaseSmartTurn):
    """End-of-turn detection with voxturn v1 (https://huggingface.co/AustinXiao/voxturn-v1)."""

    def __init__(self, *, model: str = DEFAULT_MODEL, device: str = "cuda", fast: bool = True,
                 params: VoxTurnParams | None = None, **kwargs):
        """Load the model (a Hugging Face repo id or a local copy of it) and warm it up.

        Args:
            model: Hugging Face repo id or local directory holding voxturn.py and the weights.
            device: Torch device; a CUDA GPU is expected for real-time use.
            fast: voxturn's fast mode (merged weights and CUDA graphs).
            params: Turn parameters; defaults to VoxTurnParams().
            **kwargs: Passed to BaseSmartTurn (e.g. sample_rate).
        """
        super().__init__(params=params or VoxTurnParams(), **kwargs)
        import importlib.util
        import os
        if not os.path.isdir(model):
            from huggingface_hub import snapshot_download
            model = snapshot_download(model)
        spec = importlib.util.spec_from_file_location("voxturn", os.path.join(model, "voxturn.py"))
        voxturn = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(voxturn)
        logger.debug(f"Loading voxturn from {model}...")
        self._model = voxturn.VoxTurn.from_pretrained(model, device=device, fast=fast)
        self._lock = threading.Lock()        # the analyzer's thread and the event loop share the model
        self._model.predict(np.zeros(16000, np.float32), lang=self._params.lang)   # warm-up (CUDA graph capture)
        self._next_check_ms = 0.0
        self._pending: MetricsData | None = None
        logger.debug("Loaded voxturn")

    @property
    def params(self) -> VoxTurnParams:
        return self._params

    def _predict_endpoint(self, audio_array: np.ndarray) -> dict[str, Any]:
        """P(end of turn) for the turn's audio (float32 at the pipeline's sample rate; voxturn uses the last 8 s)."""
        with self._lock:
            p = self._model.predict(audio_array, sample_rate=self.sample_rate or 16000,
                                    lang=self._params.lang)["eot"]["yes"]
        return {"prediction": 1 if p >= self._params.threshold else 0, "probability": p}

    def append_audio(self, buffer: bytes, is_speech: bool) -> EndOfTurnState:
        """BaseSmartTurn's buffering and timeout, plus the re-checks while the user stays silent."""
        state = super().append_audio(buffer, is_speech)
        recheck = self._params.recheck_secs
        if is_speech or not self._speech_triggered:
            self._next_check_ms = (recheck or 0) * 1000
        elif state == EndOfTurnState.INCOMPLETE and recheck and self._silence_ms >= self._next_check_ms:
            self._next_check_ms += recheck * 1000
            check, metrics = self._process_speech_segment(self._audio_buffer)
            if check == EndOfTurnState.COMPLETE:
                logger.debug(f"End of Turn complete by voxturn re-check after {self._silence_ms:.0f} ms of silence")
                self._pending = metrics            # handed to the strategy's follow-up analyze_end_of_turn
                self._clear(EndOfTurnState.COMPLETE)
                state = EndOfTurnState.COMPLETE
        return state

    async def analyze_end_of_turn(self) -> tuple[EndOfTurnState, MetricsData | None]:
        """As BaseSmartTurn; after a re-check ended the turn, report that re-check instead of predicting again."""
        if self._pending is not None:
            metrics, self._pending = self._pending, None
            return EndOfTurnState.COMPLETE, metrics
        return await super().analyze_end_of_turn()

    def clear(self):
        self._pending = None
        super().clear()