# -*- coding: utf-8 -*-
from typing import Generator, Literal
from urllib.error import HTTPError

import torch
from kokoro import KPipeline
from sinapsis_core.data_containers.data_packet import AudioPacket, DataContainer
from sinapsis_core.template_base.base_models import (
    OutputTypes,
    TemplateAttributes,
    TemplateAttributeType,
    UIPropertiesMetadata,
)
from sinapsis_core.template_base.template import Template
from sinapsis_core.utils.logging_utils import make_loguru

from sinapsis_kokoro.helpers.kokoro_utils import KokoroKeys, kokoro_voices
from sinapsis_kokoro.helpers.tags import Tags


class KokoroTTS(Template):
    """
    Template for text-to-speech (TTS) synthesis using the Kokoro 82M v1.0 model.
    This class handles the initialization of the TTS pipeline, speech generation,
    and packaging the output audio in the desired format.

    Usage example:

    agent:
      name: my_test_agent
    templates:
    - template_name: InputTemplate
      class_name: InputTemplate
      attributes: {}
    - template_name: KokoroTTS
      class_name: KokoroTTS
      template_input: InputTemplate
      attributes:
        speed: 1
        voice: af_heart
    """

    UIProperties = UIPropertiesMetadata(
        category="Kokoro",
        output_type=OutputTypes.AUDIO,
        tags=[Tags.AUDIO, Tags.AUDIO_GENERATION, Tags.KOKORO, Tags.SPEECH, Tags.TEXT_TO_SPEECH],
    )

    class AttributesBaseModel(TemplateAttributes):
        """
        Configuration attributes for the Kokoro TTS model.

        Args:
            speed (int | float): The speed at which the speech will be generated. Default is 1 (normal speed).
            split_pattern (str): The regular expression pattern used to split the input text into smaller chunks.
                Default is r"\n+" (split on newlines).
            voice (kokoro_voices): The voice model to use for speech synthesis. Default is "af_heart".

        Notes:
            The list of languages and voices supported by Kokoro can be found at:
            https://huggingface.co/hexgrad/Kokoro-82M/blob/main/VOICES.md
        """

        device: Literal["cpu", "cuda"] = "cpu"
        speed: int | float = 1
        split_pattern: str = r"\n+"
        voice: kokoro_voices = KokoroKeys.default_voice

    def __init__(self, attributes: TemplateAttributeType) -> None:
        """Initializes the Kokoro TTS pipeline with the provided attributes."""
        super().__init__(attributes)
        self.pipeline = self.init_pipeline()
        self.logger = make_loguru()

    def init_pipeline(self) -> KPipeline:
        """
        Initializes the Kokoro TTS pipeline with the voice model and repository id.

        Returns:
            KPipeline: The initialized TTS pipeline for generating speech.
        """
        return KPipeline(lang_code=self.attributes.voice[0], repo_id=KokoroKeys.repo_id, device=self.attributes.device)

    def _create_audio_packet(
        self,
        audio_data: torch.tensor,
        sample_rate: int,
        container: DataContainer,
    ) -> None:
        """
        Creates an audio packet from the generated audio data and adds it to the container.

        Args:
            audio_data (torch.tensor): The generated audio data (raw audio).
            sample_rate (int): The sample rate of the generated audio (typically 24000 Hz).
            container (DataContainer): The container to which the audio packet will be added.
        """
        audio_packet = AudioPacket(
            content=audio_data,
            source=self.instance_name,
            sample_rate=sample_rate,
        )
        container.audios.append(audio_packet)

    def _process_audio_chunks(self, generator: Generator, container: DataContainer) -> None:
        """
        Processes the audio chunks generated by the pipeline and creates audio packets.

        Args:
            generator: The generator that yields text, phonemes, and audio data.
            container (DataContainer): The container holding the input data.
        """
        for i, (gs, ps, audio) in enumerate(generator):
            self.logger.debug(f"Index: {i}")
            self.logger.debug(f"Text: {gs}")
            self.logger.debug(f"Phonemes: {ps}")
            if audio is not None:
                self._create_audio_packet(audio, 24000, container)
            else:
                self.logger.warning(f"Audio is None for index {i}")

    def generate_speech(self, container: DataContainer) -> None:
        """
        Generates speech from the input text in the provided data container.

        Args:
            container (DataContainer): The container holding the input text data to be converted into speech.
        """
        input_text = "".join(t.content for t in container.texts)
        generator = self.pipeline(
            input_text,
            voice=self.attributes.voice,
            speed=self.attributes.speed,
            split_pattern=self.attributes.split_pattern,
        )

        try:
            self._process_audio_chunks(generator, container)
        except HTTPError as e:
            self.logger.error(f"Unable to generate speech: {e}")

    def execute(self, container: DataContainer) -> DataContainer:
        """
        Processes the input data and generates the corresponding speech output.

        Args:
            container (DataContainer): The container holding the input text data.

        Returns:
            DataContainer: The updated container with the generated audio.
        """

        if not container.texts:
            self.logger.debug("No query to enter")
            return container

        self.generate_speech(container)

        return container

    def reset_state(self, template_name: str | None = None) -> None:
        if "cuda" in self.attributes.device:
            torch.cuda.empty_cache()
        super().reset_state(template_name)
