from __future__ import annotations

import logging
from threading import Event
from typing import Any, Iterator

import librosa
import numpy as np
import torch
from rich.console import Console
from transformers import AutoTokenizer, VitsModel

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

logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.DEBUG)

logger = logging.getLogger(__name__)

console = Console()

WHISPER_LANGUAGE_TO_FACEBOOK_LANGUAGE = {
    "en": "eng",  # English
    "fr": "fra",  # French
    "es": "spa",  # Spanish
    "ko": "kor",  # Korean
    "hi": "hin",  # Hindi
    "ar": "ara",  # Arabic
    "hy": "hyw",  # Armenian
    "az": "azb",  # Azerbaijani
    "bu": "bul",  # Bulgarian
    "ca": "cat",  # Catalan
    "nl": "nld",  # Dutch
    "fi": "fin",  # Finnish
    "de": "deu",  # German
    "el": "ell",  # Greek
    "he": "heb",  # Hebrew
    "hu": "hun",  # Hungarian
    "is": "isl",  # Icelandic
    "id": "ind",  # Indonesian
    "ka": "kan",  # Kannada
    "kk": "kaz",  # Kazakh
    "lv": "lav",  # Latvian
    "zl": "zlm",  # Malay
    "ma": "mar",  # Marathi
    "fa": "fas",  # Persian
    "po": "pol",  # Polish
    "pt": "por",  # Portuguese
    "ro": "ron",  # Romanian
    "ru": "rus",  # Russian
    "sw": "swh",  # Swahili
    "sv": "swe",  # Swedish
    "tg": "tgl",  # Tagalog
    "ta": "tam",  # Tamil
    "th": "tha",  # Thai
    "tu": "tur",  # Turkish
    "uk": "ukr",  # Ukrainian
    "ur": "urd",  # Urdu
    "vi": "vie",  # Vietnamese
    "cy": "cym",  # Welsh
}


class FacebookMMSTTSHandler(BaseHandler[TTSIn, TTSOut]):
    def setup(
        self,
        should_listen: Event,
        device: str = "cuda",
        torch_dtype: str = "float32",
        language: str = "en",
        stream: bool = True,
        chunk_size: int = 512,
        cancel_scope: CancelScope | None = None,
        speculative_turns: SpeculativeTurnTracker | None = None,
        **kwargs: Any,
    ) -> None:
        self.should_listen = should_listen
        self.cancel_scope = cancel_scope
        self.speculative_turns = speculative_turns
        self.device = device
        self.torch_dtype = getattr(torch, torch_dtype)
        self.stream = stream
        self.chunk_size = chunk_size
        self.language = language

        self._initial_language = self.language
        self.load_model(self.language)
        self.warmup()

    def load_model(self, language_code: str) -> None:
        try:
            model_name = f"facebook/mms-tts-{WHISPER_LANGUAGE_TO_FACEBOOK_LANGUAGE[language_code]}"
            logger.info(f"Loading model: {model_name}")
            self.model = VitsModel.from_pretrained(model_name).to(self.device)  # type: ignore[arg-type]
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            self.language = language_code
        except KeyError:
            logger.warning(f"Unsupported language: {language_code}. Falling back to English.")
            self.load_model("en")

    def warmup(self) -> None:
        logger.info(f"Warming up {self.__class__.__name__}")
        self.generate_audio("Hello, this is a test")

    def generate_audio(self, text: str) -> torch.Tensor | None:
        if not text:
            logger.warning("Received empty text input")
            return None

        try:
            logger.debug(f"Tokenizing text: {text}")
            logger.debug(f"Current language: {self.language}")
            logger.debug(f"Tokenizer: {self.tokenizer}")

            inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
            input_ids = inputs.input_ids.to(self.device).long()
            attention_mask = inputs.attention_mask.to(self.device)

            logger.debug(f"Input IDs shape: {input_ids.shape}, dtype: {input_ids.dtype}")
            logger.debug(f"Input IDs: {input_ids}")

            if input_ids.numel() == 0:
                logger.error("Input IDs tensor is empty")
                return None

            with torch.no_grad():
                output = self.model(input_ids=input_ids, attention_mask=attention_mask)

            logger.debug(f"Output waveform shape: {output.waveform.shape}")
            return output.waveform
        except Exception as e:
            logger.error(f"Error in generate_audio: {str(e)}")
            logger.exception("Full traceback:")
            return None

    def process(self, tts_input: TTSIn) -> Iterator[TTSOut]:
        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)

        gen = self.cancel_scope.generation if self.cancel_scope else None
        language_code = tts_input.language_code
        text = tts_input.text

        console.print(f"[green]ASSISTANT: {text}")
        logger.debug(f"Processing text: {text}")
        logger.debug(f"Language code: {language_code}")

        if language_code is not None and self.language != language_code:
            try:
                logger.info(f"Switching language from {self.language} to {language_code}")
                self.load_model(language_code)
            except KeyError:
                console.print(
                    f"[red]Language {language_code} not supported by Facebook MMS. Using {self.language} instead."
                )
                logger.warning(f"Unsupported language: {language_code}")

        audio_output = self.generate_audio(text)

        if audio_output is None or audio_output.numel() == 0:
            logger.warning("No audio output generated")
            return

        audio_numpy = audio_output.cpu().numpy().squeeze()
        logger.debug(f"Raw audio shape: {audio_numpy.shape}, dtype: {audio_numpy.dtype}")

        audio_resampled = librosa.resample(audio_numpy, orig_sr=self.model.config.sampling_rate, target_sr=16000)
        logger.debug(f"Resampled audio shape: {audio_resampled.shape}, dtype: {audio_resampled.dtype}")

        audio_int16 = (audio_resampled * 32768).astype(np.int16)
        logger.debug(f"Final audio shape: {audio_int16.shape}, dtype: {audio_int16.dtype}")

        if self.stream:
            for i in range(0, len(audio_int16), self.chunk_size):
                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_int16[i : i + self.chunk_size]
                yield np.pad(chunk, (0, self.chunk_size - len(chunk)))
        else:
            for i in range(0, len(audio_int16), self.chunk_size):
                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
                yield np.pad(
                    audio_int16[i : i + self.chunk_size],
                    (0, self.chunk_size - len(audio_int16[i : i + self.chunk_size])),
                )

    def on_session_end(self) -> None:
        if self.language != self._initial_language:
            self.load_model(self._initial_language)
        logger.debug("Facebook MMS TTS session state reset")
