"""
Custom TTS plugin using local Kokoro-82M model.

Kokoro outputs 24kHz float32 audio. We declare sample_rate=24000 and let
LiveKit handle resampling to 48kHz for WebRTC transport.
"""

import asyncio
import logging

import numpy as np
from livekit import rtc

from livekit.agents import tts
from livekit.agents.tts import AudioEmitter
from livekit.agents.types import (
    DEFAULT_API_CONNECT_OPTIONS,
    APIConnectOptions,
)

logger = logging.getLogger(__name__)


def patch_misaki_compat():
    """Patch misaki.en compatibility: MutableToken was renamed to MToken."""
    try:
        import misaki.en as en
        if not hasattr(en, "MutableToken") and hasattr(en, "MToken"):
            en.MutableToken = en.MToken
    except ImportError:
        pass


class LocalKokoroTTS(tts.TTS):
    """Text-to-Speech using local Kokoro-82M model."""

    def __init__(self, *, pipeline=None, voice: str = "zf_xiaoxiao"):
        super().__init__(
            capabilities=tts.TTSCapabilities(streaming=False),
            sample_rate=24000,  # Kokoro native output rate
            num_channels=1,
        )
        self._pipeline = pipeline  # KPipeline instance (pre-loaded in prewarm)
        self._voice = voice

    @property
    def model(self) -> str:
        return "kokoro-82m"

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

    def synthesize(
        self,
        text: str,
        *,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    ) -> "KokoroChunkedStream":
        return KokoroChunkedStream(
            tts=self,
            input_text=text,
            conn_options=conn_options,
        )


class KokoroChunkedStream(tts.ChunkedStream):
    """Non-streaming TTS synthesis with Kokoro."""

    def __init__(
        self,
        *,
        tts: LocalKokoroTTS,
        input_text: str,
        conn_options: APIConnectOptions,
    ):
        super().__init__(tts=tts, input_text=input_text, conn_options=conn_options)
        self._kokoro_tts = tts

    async def _run(self, output_emitter: AudioEmitter) -> None:
        if self._kokoro_tts._pipeline is None:
            logger.error("Kokoro pipeline not loaded")
            return

        text = self._input_text
        if not text or not text.strip():
            return

        loop = asyncio.get_event_loop()
        pipeline = self._kokoro_tts._pipeline
        voice = self._kokoro_tts._voice

        def _synthesize():
            samples = []
            for _, _, audio in pipeline(text, voice=voice):
                if hasattr(audio, "numpy"):
                    samples.append(audio.numpy())
                else:
                    samples.append(audio)
            if not samples:
                return None
            return np.concatenate(samples)

        audio_np = await loop.run_in_executor(None, _synthesize)
        if audio_np is None or len(audio_np) == 0:
            logger.warning("Kokoro produced no audio for: '%s'", text[:50])
            return

        # Convert float32 → int16 PCM
        pcm_int16 = (audio_np * 32768).clip(-32768, 32767).astype(np.int16)

        # Initialize emitter with Kokoro's native 24kHz
        request_id = text[:16].replace("\n", " ")
        output_emitter.initialize(
            request_id=request_id,
            sample_rate=24000,
            num_channels=1,
            mime_type="audio/pcm",
        )

        # Push audio data
        output_emitter.push(pcm_int16.tobytes())

        logger.info(
            "TTS synthesized: '%s' → %d samples (%.1fs)",
            text[:30],
            len(pcm_int16),
            len(pcm_int16) / 24000,
        )
