"""
Parakeet TDT Speech-to-Text Handler

Supports NVIDIA Parakeet TDT model for high-quality multilingual ASR.
- On Apple Silicon (MPS): Uses mlx-audio with mlx-community/parakeet-tdt-0.6b-v3
- On CUDA/CPU: Uses nano-parakeet (pure PyTorch) with nvidia/parakeet-tdt-0.6b-v3

Model supports 25 European languages with automatic language detection.
"""

from __future__ import annotations

import logging
from contextlib import contextmanager
from sys import platform
from threading import Lock
from time import perf_counter
from typing import Any, Iterator, Optional

import numpy as np
from rich.console import Console
from rich.text import Text

from speech_to_speech.pipeline.handler_types import STTIn, STTOut
from speech_to_speech.pipeline.messages import PartialTranscription, Transcription
from speech_to_speech.STT.base_stt_handler import BaseSTTHandler
from speech_to_speech.STT.smart_progressive_streaming import PartialTranscription as ProgressiveStreamPartial
from speech_to_speech.utils.mlx_lock import MLXLockContext

try:
    from lingua import Language, LanguageDetectorBuilder

    LINGUA_AVAILABLE = True
except ImportError:
    LINGUA_AVAILABLE = False

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

# Parakeet TDT v3 supports 25 European languages
SUPPORTED_LANGUAGES = [
    "en",
    "de",
    "fr",
    "es",
    "it",
    "pt",
    "nl",
    "pl",
    "ru",
    "uk",
    "cs",
    "sk",
    "hu",
    "ro",
    "bg",
    "hr",
    "sl",
    "sr",
    "da",
    "no",
    "sv",
    "fi",
    "et",
    "lv",
    "lt",
]

# Lingua uses "nb" (Bokmål) for Norwegian instead of "no"
_LINGUA_CODE_MAP = {"no": "nb"}

if LINGUA_AVAILABLE:
    _lingua_iso_to_code = {
        lang.iso_code_639_1.name.lower(): lang for lang in Language.all() if lang.iso_code_639_1 is not None
    }
    _lingua_languages = [
        _lingua_iso_to_code[_LINGUA_CODE_MAP.get(code, code)]
        for code in SUPPORTED_LANGUAGES
        if _LINGUA_CODE_MAP.get(code, code) in _lingua_iso_to_code
    ]

    def _build_lingua_detector():
        # Preloading can take multiple seconds on some hardware, including the
        # deployed server. Pay that cost at startup instead of on the first user
        # request, where it would look like slow STT.
        return LanguageDetectorBuilder.from_languages(*_lingua_languages).with_preloaded_language_models().build()

    _lingua_detector = _build_lingua_detector()


class ParakeetTDTSTTHandler(BaseSTTHandler):
    """
    Handles Speech-to-Text using NVIDIA Parakeet TDT model.

    On Apple Silicon (MPS): Uses mlx-audio with the MLX-converted model.
    On CUDA/CPU: Uses nano-parakeet (pure PyTorch) for NeMo-free inference.

    Parakeet TDT 0.6B v3 is a 600M parameter multilingual ASR model
    supporting 25 European languages with automatic language detection.
    """

    def setup(
        self,
        model_name: Optional[str] = None,
        device: str = "auto",
        compute_type: str = "float16",
        language: Optional[str] = None,
        gen_kwargs: dict[str, Any] = {},
        enable_live_transcription: bool = False,
        live_transcription_update_interval: float = 0.5,
    ) -> None:
        """
        Initialize the Parakeet TDT model.

        Args:
            model_name: Model identifier. Defaults are:
                - MPS: "mlx-community/parakeet-tdt-0.6b-v3"
                - CUDA/CPU: "nvidia/parakeet-tdt-0.6b-v3"
            device: Device to use ("auto", "cuda", "mps", "cpu")
            compute_type: Compute precision ("float16", "float32")
            language: Target language code (optional, model auto-detects)
            gen_kwargs: Additional generation kwargs
        """
        self.gen_kwargs = gen_kwargs
        self.start_language = language
        self.last_language = language if language else "en"
        self.enable_live_transcription = enable_live_transcription
        self.live_transcription_update_interval = live_transcription_update_interval
        self.compute_lock = Lock()
        self.sample_rate = 16000

        # 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/parakeet-tdt-0.6b-v3"
            else:
                model_name = "nvidia/parakeet-tdt-0.6b-v3"

        self.model_name = model_name
        self.compute_type = compute_type

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

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

        # Setup streaming handler if live transcription is enabled
        self.streaming_handler = None
        if self.enable_live_transcription:
            from speech_to_speech.STT.smart_progressive_streaming import (
                SmartProgressiveStreamingHandler,
            )

            self.streaming_handler = SmartProgressiveStreamingHandler(
                self.model,
                emission_interval=self.live_transcription_update_interval,
                max_window_size=15.0,
                sentence_buffer=2.0,
            )
            self.processing_final = False  # Track if we're processing final audio
            logger.info(f"Live transcription enabled for Parakeet TDT ({self.backend})")
        self._live_transcription_active = False
        self._live_turn_key: tuple[str | None, int | None] | None = None

        self.warmup()

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

            self.backend = "mlx"
            self.model = load_model(model_name)
            logger.info("MLX Audio Parakeet model loaded successfully")
        except ImportError as e:
            raise ImportError(
                "mlx-audio is required for Parakeet TDT on Apple Silicon. Install with: pip install mlx-audio"
            ) from e

    def _setup_nano_parakeet(self, model_name: str) -> None:
        """Setup for CUDA/CPU using nano-parakeet."""
        try:
            import torch
            from nano_parakeet import from_pretrained

            self.backend = "nano_parakeet"

            if self.device == "cuda" and not torch.cuda.is_available():
                logger.warning("CUDA requested but not available. Falling back to CPU for nano-parakeet.")
                self.device = "cpu"

            self.model = from_pretrained(model_name=model_name, device=self.device)

            logger.info(f"nano-parakeet model loaded successfully on {self.device}")
        except ImportError as e:
            raise ImportError(
                "nano-parakeet is required for Parakeet TDT on CUDA/CPU. Install with: pip install nano-parakeet"
            ) from e

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

        # Create 1 second of silence at 16kHz
        dummy_audio = np.zeros(16000, dtype=np.float32)

        try:
            if self.backend == "mlx":
                import mlx.core as mx

                # Convert to mx.array and call decode_chunk directly
                audio_mx = mx.array(dummy_audio, dtype=mx.float32)
                _ = self.model.decode_chunk(audio_mx, verbose=False)
            elif self.backend == "nano_parakeet":
                _ = self.model.transcribe(dummy_audio)
            else:
                _ = self.model.transcribe([dummy_audio], batch_size=1, verbose=False)

            logger.info("Model warmed up and ready")
        except Exception as e:
            logger.warning(f"Warmup failed: {e}")

    def process(self, vad_audio: STTIn) -> Iterator[STTOut]:
        """
        Process audio and generate transcription.

        Yields:
            :class:`PartialTranscription` or :class:`Transcription`
        """
        process_start_s = perf_counter()
        is_progressive = vad_audio.mode == "progressive"
        audio_input = vad_audio.audio

        # Ensure audio is float32 numpy array
        if not isinstance(audio_input, np.ndarray):
            audio_input = np.array(audio_input, dtype=np.float32)
        else:
            audio_input = audio_input.astype(np.float32)
        audio_duration_s = len(audio_input) / getattr(self, "sample_rate", 16000)
        item_age_s = self._item_age_s(vad_audio)

        self._prepare_live_transcription_turn(vad_audio.turn_id, vad_audio.turn_revision)

        # Handle progressive updates: yield tagged partial for TranscriptionNotifier
        if self.enable_live_transcription and is_progressive:
            # Ignore progressive updates if we're already processing final audio
            if self.processing_final:
                logger.debug("Skipping stale progressive update (final audio already received)")
                return

            # Try to acquire lock with short timeout - skip if busy
            lock_scope_start_s = perf_counter()
            with self._compute_lock_context(handler_name="ParakeetSTT-Progressive", timeout=0.01) as acquired:
                if acquired:
                    try:
                        inference_start_s = perf_counter()
                        progressive_text = self._show_progressive_transcription(audio_input)
                        inference_s = perf_counter() - inference_start_s
                        if inference_s >= 0.25:
                            logger.info(
                                "Parakeet progressive STT timing turn=%s rev=%s audio=%.3fs age=%.3fs "
                                "lock_scope=%.3fs inference=%.3fs chars=%d",
                                vad_audio.turn_id,
                                vad_audio.turn_revision,
                                audio_duration_s,
                                item_age_s,
                                perf_counter() - lock_scope_start_s,
                                inference_s,
                                len(progressive_text),
                            )
                        if progressive_text:
                            yield PartialTranscription(
                                text=progressive_text,
                                turn_id=vad_audio.turn_id,
                                turn_revision=vad_audio.turn_revision,
                            )
                            return
                    except Exception as e:
                        logger.debug(f"Progressive transcription failed: {e}")
                else:
                    logger.debug("Skipping progressive update (compute busy)")
            return

        # Handle final transcription (send to LLM)
        logger.info(
            "Parakeet final STT start turn=%s rev=%s audio=%.3fs age=%.3fs",
            vad_audio.turn_id,
            vad_audio.turn_revision,
            audio_duration_s,
            item_age_s,
        )
        inference_s = 0.0
        lock_scope_s = 0.0
        try:
            if self.enable_live_transcription:
                # Mark that we're processing final audio (ignore stale progressive updates)
                self.processing_final = True

            # Acquire lock with longer timeout for final transcription
            lock_scope_start_s = perf_counter()
            with self._compute_lock_context(handler_name="ParakeetSTT-Final", timeout=5.0) as acquired:
                lock_scope_s = perf_counter() - lock_scope_start_s
                if not acquired:
                    logger.error("Failed to acquire compute lock for final transcription")
                    pred_text = ""
                    language_code = self.last_language
                else:
                    inference_start_s = perf_counter()
                    if self.backend == "mlx":
                        pred_text, language_code = self._process_mlx_final(audio_input)
                    else:
                        pred_text, language_code = self._process_nano_parakeet(audio_input)
                    inference_s = perf_counter() - inference_start_s
                    lock_scope_s = perf_counter() - lock_scope_start_s

            # Validate and update language
            if language_code and language_code in SUPPORTED_LANGUAGES:
                self.last_language = language_code
            else:
                language_code = self.last_language

        except Exception as e:
            logger.error(f"Parakeet TDT inference failed: {e}")
            pred_text = ""
            language_code = self.last_language

        total_s = perf_counter() - process_start_s
        logger.info(
            "Parakeet final STT done turn=%s rev=%s total=%.3fs lock_scope=%.3fs inference=%.3fs chars=%d",
            vad_audio.turn_id,
            vad_audio.turn_revision,
            total_s,
            lock_scope_s,
            inference_s,
            len(pred_text),
        )
        logger.debug("Finished Parakeet TDT inference")
        self._clear_live_transcription_line()
        if pred_text.strip():
            console.print(f"[yellow]USER: {pred_text.strip()}")
            if language_code:
                console.print(f"[dim]Language: {language_code}[/dim]")

        # Reset per-utterance live transcription state only after final STT
        # completes. The streaming handler carries fixed sentence timing within
        # an utterance, and stale timing must not leak into the next turn.
        if self.enable_live_transcription:
            self.processing_final = False
            self._reset_live_transcription_state(clear_turn=True)

        yield Transcription(
            text=pred_text,
            language_code=language_code,
            turn_id=vad_audio.turn_id,
            turn_revision=vad_audio.turn_revision,
            speech_stopped_at_s=vad_audio.created_at_s,
        )

    @property
    def timing_log_level(self) -> int:
        return logging.INFO

    def should_log_timing(self, output: STTOut) -> bool:
        return isinstance(output, Transcription) and self.last_time > self.min_time_to_debug

    def _detect_language_from_text(self, text: str) -> Optional[str]:
        """
        Detect language from transcribed text using lingua-py.

        Args:
            text: Transcribed text string

        Returns:
            Detected language code or None if detection fails
        """
        if not LINGUA_AVAILABLE:
            logger.warning("lingua-py not available, cannot detect language from text")
            return None

        # Skip very short utterances where language ID is still too noisy.
        if not text or len(text.strip()) < 20:
            return None

        detected = _lingua_detector.detect_language_of(text)
        if detected is None:
            return None

        code = detected.iso_code_639_1.name.lower()
        # Map back lingua-specific codes to our supported codes
        return {v: k for k, v in _LINGUA_CODE_MAP.items()}.get(code, code)

    @contextmanager
    def _compute_lock_context(self, handler_name: str, timeout: float) -> Iterator[bool]:
        if self.backend == "mlx":
            with MLXLockContext(handler_name=handler_name, timeout=timeout) as acquired:
                yield acquired
            return

        lock_start_s = perf_counter()
        acquired = self.compute_lock.acquire(timeout=timeout)
        wait_s = perf_counter() - lock_start_s
        hold_start_s: float | None = None
        if acquired:
            if wait_s >= 0.25:
                logger.info("%s: compute lock acquired after %.2fs", handler_name, wait_s)
            else:
                logger.debug("%s: compute lock acquired after %.3fs", handler_name, wait_s)
            hold_start_s = perf_counter()
        else:
            logger.warning("%s: Failed to acquire compute lock after %.3fs (timeout=%s)", handler_name, wait_s, timeout)
        try:
            yield acquired
        finally:
            if acquired:
                assert hold_start_s is not None
                self.compute_lock.release()
                hold_s = perf_counter() - hold_start_s
                if hold_s >= 0.25:
                    logger.info("%s: compute lock released after holding %.2fs", handler_name, hold_s)
                else:
                    logger.debug("%s: compute lock released after holding %.3fs", handler_name, hold_s)

    def _show_progressive_transcription(self, audio_input: np.ndarray) -> str:
        """Run progressive transcription, print to console, and return the text."""
        result = self.streaming_handler.transcribe_incremental(audio_input)
        rich_text = Text()
        if result.fixed_text:
            rich_text.append("Live: ", style="dim")
            rich_text.append(result.fixed_text, style="yellow")
            if result.active_text:
                rich_text.append(" ", style="dim")

        if result.active_text:
            if not result.fixed_text:
                rich_text.append("Live: ", style="dim")
            rich_text.append(result.active_text, style="cyan dim")

        progressive_text = self._build_progressive_text(result)
        if progressive_text:
            self._print_live_transcription(rich_text, progressive_text)

        return progressive_text

    def _print_live_transcription(self, rich_text: Text, progressive_text: str) -> None:
        is_terminal = bool(getattr(console, "is_terminal", False))
        if is_terminal:
            self._write_live_control("\r\x1b[2K")
            if rich_text:
                console.print(self._truncate_live_transcription(rich_text), end="")
            else:
                fallback = Text("Live: ", style="dim")
                fallback.append(progressive_text, style="cyan dim")
                console.print(self._truncate_live_transcription(fallback), end="")
            self._write_live_control("\r")
            self._live_transcription_active = True
            return

        if rich_text:
            console.print(rich_text)
        else:
            console.print(f"[dim]Live: [/dim]{progressive_text}")

    def _clear_live_transcription_line(self) -> None:
        if not getattr(self, "_live_transcription_active", False):
            return
        self._write_live_control("\r\x1b[2K")
        self._live_transcription_active = False

    def _write_live_control(self, sequence: str) -> None:
        file = getattr(console, "file", None)
        if file is None:
            return
        file.write(sequence)
        file.flush()

    def _truncate_live_transcription(self, text: Text) -> Text:
        text = text.copy()
        width = getattr(console, "width", 80)
        try:
            max_width = max(1, int(width) - 1)
        except (TypeError, ValueError):
            max_width = 79
        text.truncate(max_width, overflow="ellipsis")
        return text

    def _prepare_live_transcription_turn(self, turn_id: str | None, turn_revision: int | None) -> None:
        if not self.enable_live_transcription:
            return
        turn_key = (turn_id, turn_revision)
        if getattr(self, "_live_turn_key", None) == turn_key:
            return
        self._reset_live_transcription_state(clear_turn=False)
        self._live_turn_key = turn_key

    def _reset_live_transcription_state(self, clear_turn: bool) -> None:
        self._clear_live_transcription_line()
        streaming_handler = getattr(self, "streaming_handler", None)
        if streaming_handler is not None:
            streaming_handler.reset()
        if clear_turn:
            self._live_turn_key = None

    def _build_progressive_text(self, result: ProgressiveStreamPartial) -> str:
        parts = []
        if result.fixed_text:
            parts.append(result.fixed_text.strip())
        if result.active_text:
            parts.append(result.active_text.strip())
        return " ".join(part for part in parts if part).strip()

    def _process_mlx_final(self, audio_input: np.ndarray) -> tuple[str, str]:
        """Process final audio using MLX backend with streaming handler."""
        # If we have fixed sentences from progressive updates, only transcribe the new part
        if (
            self.streaming_handler is not None
            and hasattr(self.streaming_handler, "fixed_sentences")
            and self.streaming_handler.fixed_sentences
        ):
            self._clear_live_transcription_line()

            # Get fixed text from previous progressive updates
            fixed_text = " ".join(self.streaming_handler.fixed_sentences).strip()

            # Calculate where fixed part ends in audio
            fixed_end_time = self.streaming_handler.fixed_end_time
            sample_rate = 16000
            fixed_end_sample = int(fixed_end_time * sample_rate)
            if fixed_end_sample > len(audio_input):
                logger.warning(
                    "Ignoring stale progressive fixed text: fixed_end_sample=%d exceeds final audio samples=%d",
                    fixed_end_sample,
                    len(audio_input),
                )
                pred_text, language_code = self._process_mlx(audio_input)
                return pred_text, language_code

            # Only transcribe the part after fixed sentences
            if fixed_end_sample < len(audio_input):
                remaining_audio = audio_input[fixed_end_sample:]

                import mlx.core as mx

                audio_mx = mx.array(remaining_audio, dtype=mx.float32)
                result = self.model.decode_chunk(audio_mx, verbose=False)

                if hasattr(result, "text"):
                    new_text = result.text.strip()
                else:
                    new_text = str(result).strip()

                # Combine fixed + new
                pred_text = f"{fixed_text} {new_text}".strip() if new_text else fixed_text
            else:
                # All audio already transcribed in progressive updates
                pred_text = fixed_text

            # Reset streaming handler for next utterance
            self.streaming_handler.reset()
        else:
            # No progressive updates, transcribe everything
            pred_text, language_code = self._process_mlx(audio_input)
            return pred_text, language_code

        # Determine language
        if self.start_language and self.start_language != "auto":
            language_code = self.start_language
        else:
            detected_lang = self._detect_language_from_text(pred_text)
            if detected_lang:
                language_code = detected_lang
            else:
                language_code = self.last_language

        return pred_text, language_code

    def _process_mlx(self, audio_input: np.ndarray) -> tuple[str, str]:
        """Process audio using MLX backend."""
        import mlx.core as mx

        # Convert numpy array to mx.array
        audio_mx = mx.array(audio_input, dtype=mx.float32)

        # Call decode_chunk directly with the audio array
        result = self.model.decode_chunk(audio_mx, verbose=False)

        # Extract text from result
        if hasattr(result, "text"):
            pred_text = result.text.strip()
        else:
            pred_text = str(result).strip()

        # Determine language:
        # 1. Use fixed language if specified by user
        # 2. Try to detect from transcribed text using langdetect
        # 3. Fall back to last known language
        if self.start_language and self.start_language != "auto":
            language_code = self.start_language
        else:
            # Detect language from transcribed text
            detected_lang = self._detect_language_from_text(pred_text)
            if detected_lang:
                language_code = detected_lang
            else:
                language_code = self.last_language

        return pred_text, language_code

    def _process_nano_parakeet(self, audio_input: np.ndarray) -> tuple[str, str]:
        """Process audio using nano-parakeet backend."""
        pred_text = self.model.transcribe(audio_input).strip()

        if self.start_language and self.start_language != "auto":
            language_code = self.start_language
        else:
            detected_lang = self._detect_language_from_text(pred_text)
            if detected_lang:
                language_code = detected_lang
            else:
                language_code = self.last_language

        return pred_text, language_code

    def cleanup(self) -> None:
        """Clean up model resources."""
        logger.info(f"Cleaning up {self.__class__.__name__}")
        if hasattr(self, "model"):
            del self.model

    def on_session_end(self) -> None:
        super().on_session_end()
        self.last_language = self.start_language if self.start_language else "en"
        if self.enable_live_transcription:
            self.processing_final = False
            self._reset_live_transcription_state(clear_turn=True)
        logger.debug("Parakeet TDT session state reset")
