"""
Custom STT plugin using local faster-whisper model.

Implements batch recognition (not streaming): LiveKit's AgentSession collects
audio during speech (using Silero VAD), and when silence is detected, passes
the complete AudioBuffer to _recognize_impl() for transcription.
"""

import asyncio
import logging

import numpy as np

from livekit.agents import stt
from livekit.agents.types import (
    DEFAULT_API_CONNECT_OPTIONS,
    NOT_GIVEN,
    APIConnectOptions,
    NotGivenOr,
)

logger = logging.getLogger(__name__)


class LocalWhisperSTT(stt.STT):
    """Speech-to-Text using local faster-whisper model."""

    def __init__(self, *, model=None, language: str = "zh"):
        super().__init__(
            capabilities=stt.STTCapabilities(
                streaming=False,
                interim_results=False,
            )
        )
        self._model = model  # WhisperModel instance (pre-loaded in prewarm)
        self._language = language

    @property
    def model(self) -> str:
        return "faster-whisper"

    @property
    def provider(self) -> str:
        return "local"

    async def _recognize_impl(
        self,
        buffer,
        *,
        language: NotGivenOr[str] = NOT_GIVEN,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    ) -> stt.SpeechEvent:
        if self._model is None:
            logger.error("Whisper model not loaded")
            return stt.SpeechEvent(
                type=stt.SpeechEventType.FINAL_TRANSCRIPT,
                alternatives=[stt.SpeechData(language=self._language, text="")],
            )

        # Convert AudioBuffer to numpy float32 array
        # LiveKit AudioBuffer provides audio frames; we need to collect them
        audio_frames = buffer if isinstance(buffer, list) else [buffer]
        pcm_data = bytearray()
        sample_rate = 16000

        for frame in audio_frames:
            if hasattr(frame, "data"):
                pcm_data.extend(frame.data)
                if hasattr(frame, "sample_rate"):
                    sample_rate = frame.sample_rate
            elif isinstance(frame, bytes):
                pcm_data.extend(frame)
            else:
                # Try to get raw data from the frame
                try:
                    pcm_data.extend(bytes(frame))
                except TypeError:
                    pass

        if not pcm_data:
            return stt.SpeechEvent(
                type=stt.SpeechEventType.FINAL_TRANSCRIPT,
                alternatives=[stt.SpeechData(language=self._language, text="")],
            )

        # Convert PCM int16 to float32 normalized [-1.0, 1.0]
        audio_np = np.frombuffer(bytes(pcm_data), dtype=np.int16).astype(np.float32) / 32768.0

        lang = language if language is not NOT_GIVEN else self._language

        # Run transcription in thread pool to avoid blocking
        loop = asyncio.get_event_loop()

        def _transcribe():
            segments, info = self._model.transcribe(
                audio_np,
                language=lang,
                beam_size=5,
                vad_filter=True,
            )
            return "".join(seg.text for seg in segments)

        text = await loop.run_in_executor(None, _transcribe)
        text = text.strip()

        logger.info("STT result: '%s'", text[:100] if text else "(empty)")

        return stt.SpeechEvent(
            type=stt.SpeechEventType.FINAL_TRANSCRIPT,
            alternatives=[
                stt.SpeechData(
                    language=lang,
                    text=text,
                    confidence=1.0,
                )
            ],
        )
