"""
Kokoro TTS Handler

Supports NVIDIA Kokoro TTS model for high-quality multilingual speech synthesis.
- On Apple Silicon (MPS): Uses mlx-audio with mlx-community/Kokoro-82M-bf16
- On CUDA/CPU: Uses native kokoro library with hexgrad/Kokoro-82M

Model supports 8 languages with multiple voices per language.
"""

from __future__ import annotations

import logging
from sys import platform
from threading import Event
from typing import Any, Iterator, Optional

import numpy as np
from rich.console import Console

from speech_to_speech.baseHandler import BaseHandler
from speech_to_speech.pipeline.cancel_scope import CancelScope
from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse
from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker
from speech_to_speech.utils.mlx_lock import MLXLockContext

logger = logging.getLogger(__name__)
console = Console()

# Language code mapping from Whisper/langdetect language codes to Kokoro lang codes
WHISPER_LANGUAGE_TO_KOKORO_LANG = {
    "en": "b",  # British English
    "ja": "j",  # Japanese
    "zh": "z",  # Chinese
    "fr": "f",  # French
    "es": "e",  # Spanish
    "it": "i",  # Italian
    "pt": "p",  # Portuguese
    "hi": "h",  # Hindi
    # Additional European languages from Parakeet - map to closest Kokoro language
    "de": "b",  # German -> British English (no German voice available)
    "nl": "b",  # Dutch -> British English
    "pl": "b",  # Polish -> British English
    "ru": "b",  # Russian -> British English
    "uk": "b",  # Ukrainian -> British English
}

# Default voices for each Kokoro language code
# These are native voices that sound natural for each language
# Voice naming: first letter = language, second letter = gender (f=female, m=male)
# Available voices per language:
#   a (American): af_alloy, af_aoede, af_bella, af_heart, af_jessica, af_kore, af_nicole, af_nova, af_river, af_sarah, af_sky
#                 am_adam, am_echo, am_eric, am_fenrir, am_liam, am_michael, am_onyx, am_puck, am_santa
#   b (British):  bf_alice, bf_emma, bf_isabella, bf_lily, bm_daniel, bm_fable, bm_george, bm_lewis
#   e (Spanish):  ef_dora, em_alex, em_santa
#   f (French):   ff_siwis
#   h (Hindi):    hf_alpha, hf_beta, hm_omega, hm_psi
#   i (Italian):  if_sara, im_nicola
#   j (Japanese): jf_alpha, jf_gongitsune, jf_nezumi, jf_tebukuro, jm_kumo
#   p (Portuguese): pf_dora, pm_alex, pm_santa
#   z (Chinese):  zf_xiaobei, zf_xiaoni, zf_xiaoxiao, zf_xiaoyi, zm_yunjian, zm_yunxi, zm_yunxia, zm_yunyang
KOKORO_LANG_DEFAULT_VOICES = {
    "a": "af_heart",  # American English female
    "b": "bm_fable",  # British English male
    "e": "ef_dora",  # Spanish female
    "f": "ff_siwis",  # French female
    "h": "hf_alpha",  # Hindi female
    "i": "if_sara",  # Italian female
    "j": "jf_alpha",  # Japanese female
    "p": "pf_dora",  # Portuguese female
    "z": "zf_xiaobei",  # Chinese female
}


class KokoroTTSHandler(BaseHandler[TTSIn, TTSOut]):
    """
    Handles Text-to-Speech using Kokoro TTS model.

    On Apple Silicon (MPS): Uses mlx-audio with the MLX-converted model.
    On CUDA/CPU: Uses native kokoro library for optimal performance.

    Kokoro 82M is a 82M parameter multilingual TTS model
    supporting 8 languages with multiple voices per language.
    """

    def setup(
        self,
        should_listen: Event,
        model_name: Optional[str] = None,
        device: str = "auto",
        voice: str = "bm_fable",
        lang_code: str = "b",
        speed: float = 1.0,
        blocksize: int = 512,
        gen_kwargs: dict[str, Any] | None = None,
        cancel_scope: CancelScope | None = None,
        speculative_turns: SpeculativeTurnTracker | None = None,
    ) -> None:
        """
        Initialize the Kokoro TTS model.

        Args:
            model_name: Model identifier. Defaults are:
                - MPS: "mlx-community/Kokoro-82M-bf16"
                - CUDA/CPU: "hexgrad/Kokoro-82M"
            device: Device to use ("auto", "cuda", "mps", "cpu")
            voice: Voice identifier (e.g. "bm_fable")
            lang_code: Language code (e.g. "b" for British English)
            speed: Speech speed multiplier
            blocksize: Audio chunk size for streaming
            gen_kwargs: Unused, for pipeline compatibility
        """
        self.should_listen = should_listen
        self.voice = voice
        self.lang_code = lang_code
        self.speed = speed
        self.blocksize = blocksize
        self.cancel_scope = cancel_scope
        self.speculative_turns = speculative_turns

        # Determine device
        if device == "auto":
            if platform == "darwin":
                self.device = "mps"
            else:
                import torch

                self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        # Set default model based on device
        if model_name is None:
            if self.device == "mps":
                model_name = "mlx-community/Kokoro-82M-bf16"
            else:
                model_name = "hexgrad/Kokoro-82M"

        self.model_name = model_name

        logger.info(f"Loading Kokoro model: {model_name} on {self.device}")

        if self.device == "mps":
            self._setup_mlx(model_name)
        else:
            self._setup_kokoro(model_name)

        self._initial_voice = self.voice
        self._initial_lang_code = self.lang_code

        self.warmup()

    def _setup_mlx(self, model_name: str) -> None:
        """Setup for Apple Silicon using mlx-audio."""
        try:
            from mlx_audio.tts.utils import load_model

            self.backend = "mlx"
            self.model = load_model(model_name)
            logger.info("MLX Audio Kokoro model loaded successfully")

            # Get or create the pipeline for our language and preload the voice
            # This avoids the voice being reloaded on every generate() call
            self._pipeline = self.model._get_pipeline(self.lang_code)
            self._voice_tensor = self._pipeline.load_voice(self.voice)
            logger.info(f"Preloaded voice: {self.voice}")

            # Preload voices for common languages to avoid download delays during inference
            self._preload_multilingual_voices()
        except ImportError as e:
            message = str(e)
            if "misaki" in message or "espeakng_loader" in message:
                raise ImportError(
                    "Kokoro TTS on Apple Silicon requires additional mlx-audio TTS dependencies. "
                    f"Missing dependency: {message}. "
                    "Install with: pip install misaki espeakng-loader"
                ) from e
            raise ImportError(
                "mlx-audio is required for Kokoro TTS on Apple Silicon. Install with: pip install mlx-audio"
            ) from e

    def _setup_kokoro(self, model_name: str) -> None:
        """Setup for CUDA/CPU using native kokoro library."""
        try:
            from kokoro import KPipeline

            self.backend = "kokoro"
            self.pipeline = KPipeline(lang_code=self.lang_code)
            logger.info("Native Kokoro pipeline loaded successfully")
        except ImportError as e:
            raise ImportError(
                "kokoro is required for Kokoro TTS on CUDA/CPU. "
                "Install with: pip install kokoro>=0.9.2 soundfile\n"
                "Also ensure espeak-ng is installed: apt-get install espeak-ng (Linux) or brew install espeak-ng (macOS)"
            ) from e

    def _preload_multilingual_voices(self) -> None:
        """Preload voices for common languages to avoid download delays during inference."""
        # Only preload a few commonly used language voices to avoid excessive startup time
        # Users speaking other languages will experience a one-time download delay
        # preload_langs = ["a", "b", "e", "f", "h", "i", "j", "p", "z"] # English, Spanish, French, Hindi, Italian, Japanese, Portuguese
        preload_langs = ["a", "e", "f"]  # English, Spanish, French

        for lang_code in preload_langs:
            voice = KOKORO_LANG_DEFAULT_VOICES.get(lang_code)
            if voice:
                try:
                    logger.info(f"Preloading voice for language '{lang_code}': {voice}")
                    pipeline = self.model._get_pipeline(lang_code)
                    pipeline.load_voice(voice)
                except Exception as e:
                    logger.warning(f"Failed to preload voice {voice}: {e}")

    def warmup(self) -> None:
        """Warm up the model with a dummy inference."""
        logger.info(f"Warming up {self.__class__.__name__}")

        # Run a short dummy inference to warm up the model
        if self.backend == "mlx":
            for _ in self.model.generate(
                text="Hello",
                voice=self.voice,
                speed=self.speed,
                lang_code=self.lang_code,
            ):
                pass
        else:
            for _ in self.pipeline("Hello", voice=self.voice, speed=self.speed):
                pass

        logger.info(f"{self.__class__.__name__} warmed up")

    def process(self, tts_input: TTSIn) -> Iterator[TTSOut]:
        """
        Process text input and generate audio output.

        Yields:
            Audio chunks as numpy int16 arrays
        """
        speculative_turns = getattr(self, "speculative_turns", None)
        if isinstance(tts_input, EndOfResponse):
            if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
                tts_input.turn_id,
                tts_input.turn_revision,
            ):
                return
            yield AUDIO_RESPONSE_DONE
            return

        if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
            tts_input.turn_id,
            tts_input.turn_revision,
        ):
            logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision)
            return
        if speculative_turns:
            speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision)

        runtime_config = tts_input.runtime_config
        response = tts_input.response
        language_code = tts_input.language_code
        text = tts_input.text

        voice: Optional[str] = None
        if response and response.audio and response.audio.output and response.audio.output.voice:
            voice = str(response.audio.output.voice)
        if not voice and runtime_config:
            audio_cfg = runtime_config.session.audio
            audio_output = audio_cfg.output if audio_cfg is not None else None
            voice = str(audio_output.voice) if audio_output is not None and audio_output.voice else None
        if voice:
            self.voice = voice

        if self.backend == "mlx":
            yield from self._process_mlx(text, language_code)
        else:
            yield from self._process_kokoro(text, language_code)

    def _process_mlx(self, llm_sentence: str, language_code: Optional[str] = None) -> Iterator[np.ndarray]:
        """Process using MLX backend with Apple Silicon optimizations."""
        from scipy.signal import resample_poly

        gen = self.cancel_scope.generation if self.cancel_scope else None
        with MLXLockContext(handler_name="KokoroTTS", timeout=10.0):
            if language_code is not None:
                new_lang_code = WHISPER_LANGUAGE_TO_KOKORO_LANG.get(language_code, self.lang_code)
                if new_lang_code != self.lang_code:
                    new_voice = KOKORO_LANG_DEFAULT_VOICES.get(new_lang_code, self.voice)
                    logger.info(
                        f"Language change detected: {self.lang_code} -> {new_lang_code}, voice: {self.voice} -> {new_voice}"
                    )
                    try:
                        new_pipeline = self.model._get_pipeline(new_lang_code)
                        new_voice_tensor = new_pipeline.load_voice(new_voice)
                        self.lang_code = new_lang_code
                        self.voice = new_voice
                        self._pipeline = new_pipeline
                        self._voice_tensor = new_voice_tensor
                    except Exception as e:
                        logger.warning(
                            f"Failed to switch language/voice: {e}. Keeping current language: {self.lang_code}"
                        )

            console.print(f"[green]ASSISTANT: {llm_sentence}")

            # Generate audio using the preloaded pipeline directly
            # This avoids the voice reload that happens in model.generate()
            for result in self._pipeline(
                text=llm_sentence,
                voice=self.voice,
                speed=self.speed,
            ):
                if result.audio is None:
                    continue
                # result.audio is an mx.array with shape (1, samples), convert to numpy and squeeze
                audio = np.array(result.audio, dtype=np.float32).squeeze(0)

                # Trim silence from start and end of audio
                # Kokoro generates ~250ms of silence at the start and variable silence at the end
                threshold = 0.01
                abs_audio = np.abs(audio)
                # Find first and last samples above threshold
                above_threshold = abs_audio > threshold
                if np.any(above_threshold):
                    start_idx = int(np.argmax(above_threshold))
                    end_idx = len(audio) - int(np.argmax(above_threshold[::-1]))
                    # Add small padding (5ms = 120 samples at 24kHz) to avoid cutting speech
                    padding = 120
                    start_idx = max(0, start_idx - padding)
                    end_idx = min(len(audio), end_idx + padding)
                    audio = audio[start_idx:end_idx]

                # Kokoro outputs at 24kHz, resample to 16kHz for the pipeline
                # Using scipy's polyphase resampling (fast and high quality)
                # 16000/24000 = 2/3, so up=2, down=3
                audio = resample_poly(audio, up=2, down=3)

                # Convert to int16 format
                audio = (audio * 32768).astype(np.int16)

                # Yield audio in fixed-size chunks
                for i in range(0, len(audio), self.blocksize):
                    if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen):
                        logger.info("TTS generation cancelled (interruption)")
                        return
                    chunk = audio[i : i + self.blocksize]
                    # Pad the last chunk if necessary
                    if len(chunk) < self.blocksize:
                        chunk = np.pad(chunk, (0, self.blocksize - len(chunk)))
                    logger.debug(f"TTS yielding audio chunk: {len(chunk)} samples")
                    yield chunk

    def _process_kokoro(self, llm_sentence: str, language_code: Optional[str] = None) -> Iterator[np.ndarray]:
        """Process using native kokoro library."""
        from scipy.signal import resample_poly

        gen = self.cancel_scope.generation if self.cancel_scope else None
        if language_code is not None:
            new_lang_code = WHISPER_LANGUAGE_TO_KOKORO_LANG.get(language_code, self.lang_code)
            if new_lang_code != self.lang_code:
                new_voice = KOKORO_LANG_DEFAULT_VOICES.get(new_lang_code, self.voice)
                logger.info(
                    f"Language change detected: {self.lang_code} -> {new_lang_code}, voice: {self.voice} -> {new_voice}"
                )
                self.lang_code = new_lang_code
                self.voice = new_voice
                from kokoro import KPipeline

                self.pipeline = KPipeline(lang_code=self.lang_code)

        console.print(f"[green]ASSISTANT: {llm_sentence}")

        # Generate audio using Kokoro
        # The pipeline yields tuples of (graphemes, phonemes, audio)
        for gs, ps, audio in self.pipeline(llm_sentence, voice=self.voice, speed=self.speed):
            if audio is None:
                continue

            # Ensure audio is numpy array
            if not isinstance(audio, np.ndarray):
                audio = np.array(audio, dtype=np.float32)
            else:
                audio = audio.astype(np.float32)

            # Kokoro outputs at 24kHz, resample to 16kHz for the pipeline
            # Using scipy's polyphase resampling (fast and high quality)
            # 16000/24000 = 2/3, so up=2, down=3
            audio = resample_poly(audio, up=2, down=3)

            # Convert to int16 format
            audio = (audio * 32768).astype(np.int16)

            # Yield audio in fixed-size chunks
            for i in range(0, len(audio), self.blocksize):
                if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen):
                    logger.info("TTS generation cancelled (interruption)")
                    return
                chunk = audio[i : i + self.blocksize]
                # Pad the last chunk if necessary
                if len(chunk) < self.blocksize:
                    chunk = np.pad(chunk, (0, self.blocksize - len(chunk)))
                yield chunk

    def on_session_end(self) -> None:
        self.voice = self._initial_voice
        self.lang_code = self._initial_lang_code
        if self.backend == "mlx":
            try:
                self._pipeline = self.model._get_pipeline(self.lang_code)
                self._voice_tensor = self._pipeline.load_voice(self.voice)
            except Exception as e:
                logger.warning(f"Failed to restore initial voice/language on session end: {e}")
        else:
            from kokoro import KPipeline

            self.pipeline = KPipeline(lang_code=self.lang_code)
        logger.debug("Kokoro TTS session state reset")
