# Refactored imports for better readability and grouping
import os
import time
from queue import Queue
from typing import Optional, Union, ClassVar, Type, Dict, Any

import torch
import datasets
from datasets import (
    load_dataset, 
    DatasetDict, 
    Dataset, 
    IterableDatasetDict, 
    IterableDataset
)
from transformers import (
    AutoTokenizer, PreTrainedModel, ProcessorMixin,
    SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
)

from airunner.handlers.tts.tts_model_manager import TTSModelManager
from airunner.enums import ModelType, ModelStatus, SpeechT5Voices
from airunner.utils.memory import clear_memory
from airunner.settings import AIRUNNER_LOCAL_FILES_ONLY

class SpeechT5ModelManager(TTSModelManager):
    """
    SpeechT5 implementation of the TTSModelManager.
    Uses the SpeechT5ForTextToSpeech model and SpeechT5Processor.
    """
    target_model: ClassVar[str] = "microsoft/speecht5_tts"
    model_class: ClassVar[Type[PreTrainedModel]] = SpeechT5ForTextToSpeech
    processor_class: ClassVar[Type[ProcessorMixin]] = SpeechT5Processor
    vocoder_class: ClassVar[Type[PreTrainedModel]] = SpeechT5HifiGan
    tokenizer_class: ClassVar[Type[AutoTokenizer]] = AutoTokenizer
    dataset_path: ClassVar[str] = "Matthijs/cmu-arctic-xvectors"
    speakers: Dict[str, str] = {
        SpeechT5Voices.US_MALE.value: "bdl",
        SpeechT5Voices.US_MALE_2.value: "rms",
        SpeechT5Voices.US_FEMALE.value: "slt",
        SpeechT5Voices.US_FEMALE_2.value: "clb",
        SpeechT5Voices.CANADIAN_MALE.value: "jmk",
        SpeechT5Voices.SCOTTISH_MALE.value: "awb",
        SpeechT5Voices.INDIAN_MALE.value: "ksp",
    }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.vocoder = None
        self._character_replacement_map = {
            "\n": " ",
            "’": "'",
            "-": " "
        }
        self._model: Optional[Type[PreTrainedModel]] = None
        self._vocoder: Optional[Type[PreTrainedModel]] = None
        self._processor: Optional[Type[ProcessorMixin]] = None
        self._tokenizer: Optional[Type[AutoTokenizer]] = None
        self._text_queue = Queue()
        self._speaker_embeddings = None
        self._dataset = None
        self._sentences = []
        self._do_interrupt = False
        self._cancel_generated_speech = False
        self._paused = False

    # Refactored properties for better readability
    @property
    def model(self) -> Type[PreTrainedModel]:
        return self._model

    @model.setter
    def model(self, value: Type[PreTrainedModel]):
        self._model = value

    @property
    def processor_path(self) -> str:
        return self.tts_path(self.speech_t5_settings.processor_path)

    @property
    def model_path(self) -> str:
        return self.tts_path(self.speech_t5_settings.model_path)

    @property
    def vocoder_path(self) -> str:
        return self.tts_path(self.speech_t5_settings.vocoder_path)

    @property
    def dtype(self) -> torch.dtype:
        return torch.float16

    @property
    def vocoder(self) -> Optional[Type[PreTrainedModel]]:
        return self._vocoder

    @vocoder.setter
    def vocoder(self, value: Optional[Type[PreTrainedModel]]):
        self._vocoder = value

    @property
    def processor(self) -> Optional[Type[ProcessorMixin]]:
        return self._processor

    @processor.setter
    def processor(self, value: Optional[Type[ProcessorMixin]]):
        self._processor = value

    @property
    def tokenizer(self) -> Optional[Type[AutoTokenizer]]:
        return self._tokenizer

    @tokenizer.setter
    def tokenizer(self, value: Optional[Type[AutoTokenizer]]):
        self._tokenizer = value

    @property
    def device(self) -> torch.device:
        """Return the appropriate device based on settings."""
        return torch.device("cuda" if self.tts_settings.use_cuda else "cpu")

    @property
    def torch_dtype(self) -> torch.dtype:
        """Return the consistent torch dtype to use across models and inputs."""
        return torch.float16 if self.tts_settings.use_cuda else torch.float32

    def tts_path(self, path: str) -> str:
        return os.path.join(self.path_settings.tts_model_path, path)

    # Refactored status methods for consistency
    def _set_status_unloaded(self):
        self.change_model_status(ModelType.TTS, ModelStatus.UNLOADED)

    def _set_status_loading(self):
        self.change_model_status(ModelType.TTS, ModelStatus.LOADING)

    def _set_status_loaded(self):
        self.change_model_status(ModelType.TTS, ModelStatus.LOADED)

    def _set_status_failed(self):
        self.change_model_status(ModelType.TTS, ModelStatus.FAILED)

    # Refactored load/unload methods for better error handling
    def load(self, target_model=None):
        if self.model_status is ModelStatus.LOADING:
            return
        if self.model_status in (
            ModelStatus.LOADED,
            ModelStatus.READY,
            ModelStatus.FAILED
        ):
            self.unload()
        self.logger.debug(f"Loading text-to-speech")
        self._set_status_loading()
        self._load_model()
        self._load_vocoder()
        self._load_processor()
        self._load_speaker_embeddings()
        self._load_tokenizer()
        if (
            self.model is not None
            and self.vocoder is not None
            and self.processor is not None
            and self._speaker_embeddings is not None
            and self.tokenizer is not None
        ):
            self._set_status_loaded()
        else:
            self._set_status_failed()

    def unload(self):
        if self.model_status is ModelStatus.LOADING:
            return
        self._set_status_loading()
        self.model = None
        self.processor = None
        self.vocoder = None
        self._unload_speaker_embeddings()
        self.tokenizer = None
        clear_memory(self.memory_settings.default_gpu_tts)
        self._set_status_unloaded()

    # Refactored generate method for better exception handling
    def generate(self, message: str):
        if self.model_status is not ModelStatus.LOADED:
            return None
        if self._do_interrupt or self._paused:
            return None
        try:
            return self._do_generate(message)
        except torch.cuda.OutOfMemoryError:
            self.logger.error("Out of memory")
            return None

    def _load_model(self):
        if self.model_class is None:
            return
        self.logger.debug(f"Loading model {self.model_path}")
        try:
            self.model = self.model_class.from_pretrained(
                self.model_path,
                local_files_only=AIRUNNER_LOCAL_FILES_ONLY,
                torch_dtype=self.torch_dtype,
                device_map=self.device
            )
        except EnvironmentError as _e:
            self.logger.error(f"Failed to load model {_e}")

    def _load_tokenizer(self):
        self.logger.debug("Loading tokenizer")

        try:
            self.tokenizer = self.tokenizer_class.from_pretrained(
                self.model_path,
                device_map=self.device,
                torch_dtype=self.torch_dtype,
                local_files_only=AIRUNNER_LOCAL_FILES_ONLY,
                trust_remote_code=False
            )
        except Exception as e:
            self.logger.error("Failed to load tokenizer")
            self.logger.error(e)

    def _load_vocoder(self):
        self.logger.debug(f"Loading Vocoder {self.vocoder_path}")
        try:
            self.vocoder = self.vocoder_class.from_pretrained(
                self.vocoder_path,
                local_files_only=AIRUNNER_LOCAL_FILES_ONLY,
                torch_dtype=self.torch_dtype,
                device_map=self.device
            )
        except Exception as e:
            self.logger.error("Failed to load vocoder")
            self.logger.error(e)

    def _load_processor(self):
        self.logger.debug("Loading Procesor")
        if self.processor_class:
            try:
                self.processor = self.processor_class.from_pretrained(
                    self.processor_path,
                    local_files_only=AIRUNNER_LOCAL_FILES_ONLY,
                    torch_dtype=self.torch_dtype,
                    device_map=self.device
                )
            except Exception as e:
                self.logger.error("Failed to load processor")
                self.logger.error(e)

    def _load_speaker_embeddings(self):
        self.logger.debug("Loading speaker embeddings...")
        embeddings_dataset = load_dataset(
            self.dataset_path, 
            split=datasets.Split.VALIDATION
        )
        speaker_key = self.speakers[self.speech_t5_settings.voice]
        embeddings = self._load_dataset_by_speaker_key(
            speaker_key,
            embeddings_dataset
        )
        if embeddings is None:
            self.logger.error(
                "Failed to load speaker embeddings. Using fallback"
            )
            embeddings = self._load_dataset_by_speaker_key(
                "slt",
                embeddings_dataset
            )

        self._speaker_embeddings = embeddings
        if self._speaker_embeddings is not None:
            if self.use_cuda:
                self.logger.info("Moving speaker embeddings to CUDA")
                self._speaker_embeddings = self._speaker_embeddings.to(
                    torch.bfloat16
                ).cuda()
        else:
            self.logger.error("Failed to load speaker embeddings")
        
    @staticmethod
    def _extract_speaker_key(filename):
        return filename.split("_")[2]

    def _load_dataset_by_speaker_key(
        self, 
        speaker_key: str,
        embeddings_dataset: Union[
            DatasetDict, 
            Dataset, 
            IterableDatasetDict, 
            IterableDataset
        ] = None
    ) -> Optional[torch.Tensor]:
        embeddings = None
        for entry in embeddings_dataset:
            speaker = self._extract_speaker_key(entry["filename"])
            if speaker_key == speaker:
                embeddings = torch.tensor(entry["xvector"]).unsqueeze(0)
                break
        return embeddings

    def reload_speaker_embeddings(self):
        self._unload_speaker_embeddings()
        self._load_speaker_embeddings()

    def _unload_speaker_embeddings(self):
        self._speaker_embeddings = None
        clear_memory(self.memory_settings.default_gpu_tts)

    def _do_generate(self, message: str):
        self.logger.debug("Generating text-to-speech with T5")
        text = self._prepare_text(message)

        if text == "":
            return None

        self.logger.debug("Processing inputs...")

        inputs = self.processor(
            text=text,
            return_tensors="pt",
        )
        inputs = self._move_inputs_to_device(inputs)

        self.logger.debug("Generating speech...")
        start = time.time()
        
        # Use consistent device and dtype handling
        try:
            speaker_embeddings = self._speaker_embeddings.to(self.torch_dtype).to(self.device)
            vocoder = self.vocoder.to(self.torch_dtype).to(self.device)
            
            speech = self.model.generate(
                **inputs,
                speaker_embeddings=speaker_embeddings,
                vocoder=vocoder,
                max_length=100
            )
        except Exception as e:
            self.logger.error(f"Failed to generate speech: {str(e)}")
            self._cancel_generated_speech = False
            return None

        if not self._cancel_generated_speech:
            self.logger.debug(f"Generated speech in {time.time() - start:.2f} seconds")
            response = speech.cpu().float().numpy()
            return response
        if not self._do_interrupt:
            self.logger.debug(f"Skipping generated speech: {text}")
            self._cancel_generated_speech = False
        return None

    def _move_inputs_to_device(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Move input tensors to the appropriate device."""
        try:
            device = self.device
            for key in inputs:
                if isinstance(inputs[key], torch.Tensor):
                    inputs[key] = inputs[key].to(device)
                elif isinstance(inputs[key], dict):
                    for subkey in inputs[key]:
                        if isinstance(inputs[key][subkey], torch.Tensor):
                            inputs[key][subkey] = inputs[key][subkey].to(device)
                            
        except Exception as e:
            self.logger.error(f"Failed to move inputs to device: {str(e)}")
            
        return inputs

    def unblock_tts_generator_signal(self):
        self.logger.debug("Unblocking text-to-speech generation...")
        self._do_interrupt = False
        self._paused = False

    def interrupt_process_signal(self):
        self._do_interrupt = True
        self._cancel_generated_speech = False
        self._paused = True
        self._text_queue = Queue()

    def cleanup(self):
        """Clean up resources when the handler is no longer needed."""
        self.unload()
        self._text_queue = Queue()
        self._sentences = []
