from __future__ import annotations

import logging
from typing import Any, Iterator, Optional

import numpy as np
from rich.console import Console

from speech_to_speech.pipeline.handler_types import STTIn, STTOut
from speech_to_speech.pipeline.messages import Transcription
from speech_to_speech.STT.base_stt_handler import BaseSTTHandler

logger = logging.getLogger(__name__)

console = Console()

SUPPORTED_LANGUAGES = [
    "en",
    "fr",
    "es",
    "zh",
    "ja",
    "ko",
    "hi",
    "de",
    "pt",
    "pl",
    "it",
    "nl",
]


class MLXAudioWhisperSTTHandler(BaseSTTHandler):
    """
    Handles the Speech To Text generation using MLX Audio's Whisper implementation.
    Optimized for Apple Silicon using the MLX framework.
    """

    def setup(
        self,
        model_name: str = "mlx-community/whisper-large-v3-turbo",
        language: Optional[str] = None,
        gen_kwargs: dict[str, Any] = {},
    ) -> None:
        from mlx_audio.stt.generate import load_model
        from transformers import WhisperProcessor

        self.model_name = model_name
        self.start_language = language
        self.last_language = language
        self.gen_kwargs = gen_kwargs

        # Load the model directly
        logger.info(f"Loading model {model_name}...")
        self.model = load_model(model_name)

        # Check if processor was loaded, if not, load it manually from original model
        if self.model._processor is None:
            logger.info("Processor not found in MLX model, loading from original Whisper model...")
            # Map MLX model names to their original Whisper counterparts
            processor_model_map = {
                "mlx-community/whisper-large-v3-turbo": "openai/whisper-large-v3",
                "mlx-community/whisper-large-v3": "openai/whisper-large-v3",
                "mlx-community/whisper-medium": "openai/whisper-medium",
                "mlx-community/whisper-small": "openai/whisper-small",
                "mlx-community/whisper-base": "openai/whisper-base",
                "mlx-community/whisper-tiny": "openai/whisper-tiny",
            }

            # Get the appropriate processor model name
            processor_model = processor_model_map.get(model_name, "openai/whisper-large-v3")
            logger.info(f"Loading processor from {processor_model}...")

            try:
                self.model._processor = WhisperProcessor.from_pretrained(processor_model)
                logger.info("Processor loaded successfully")
            except Exception as e:
                logger.error(f"Failed to load processor: {e}")
                raise

        logger.info(f"Model {model_name} loaded successfully")

        self.warmup()

    def warmup(self) -> None:
        logger.info(f"Warming up {self.__class__.__name__}")

        # Warmup with a dummy input
        dummy_audio = np.zeros(16000, dtype=np.float32)

        try:
            # Pre-warm the model by running a transcription
            _ = self.model.generate(dummy_audio, 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]:
        logger.debug("inferring mlx-audio whisper...")

        assert isinstance(vad_audio.audio, np.ndarray), "Audio must be a numpy array"
        audio_input = vad_audio.audio.astype(np.float32)

        # Prepare generation kwargs - only pass valid parameters
        gen_kwargs = {}

        # Add language if specified
        if self.start_language and self.start_language != "auto":
            gen_kwargs["language"] = self.start_language

        try:
            # Generate transcription directly using model.generate
            result = self.model.generate(audio_input, verbose=False, **gen_kwargs)

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

            # Try to detect language from result if available
            if hasattr(result, "language"):
                language_code = result.language
            elif self.start_language and self.start_language != "auto":
                language_code = self.start_language
            else:
                # Default to last known language or English
                language_code = self.last_language if self.last_language else "en"

            # Validate language code
            if language_code not in SUPPORTED_LANGUAGES:
                logger.warning(f"Detected unsupported language: {language_code}")
                if self.last_language in SUPPORTED_LANGUAGES:
                    language_code = self.last_language
                else:
                    language_code = "en"
            else:
                self.last_language = language_code

        except Exception as e:
            logger.error(f"MLX Audio Whisper inference failed: {e}")
            pred_text = ""
            language_code = self.last_language if self.last_language else "en"

        logger.debug("finished mlx-audio whisper inference")
        console.print(f"[yellow]USER: {pred_text}")
        logger.debug(f"Language Code: {language_code}")

        if self.start_language == "auto":
            language_code += "-auto"

        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,
        )