import asyncio
import json
import re
import time
from collections.abc import AsyncGenerator

import numpy as np
import onnxruntime as rt
from numpy.typing import NDArray

from .config import MAX_PHONEME_LENGTH, SAMPLE_RATE, EspeakConfig, KoKoroConfig
from .log import log
from .tokenizer import Tokenizer
from .trim import trim as trim_audio


class Kokoro:
    @classmethod
    def from_session(
        cls,
        session: rt.InferenceSession,
        voices_path: str,
        espeak_config: EspeakConfig | None = None,
        vocab_config: dict | str | None = None,
    ):
        instance = cls.__new__(cls)
        instance.sess = session
        model_path = getattr(session, "_model_path", getattr(session, "_path", ""))
        instance.config = KoKoroConfig(model_path, voices_path, espeak_config)
        instance.config.validate()
        instance.voices = np.load(voices_path)
        vocab = instance._load_vocab(vocab_config)
        instance.tokenizer = Tokenizer(espeak_config, vocab=vocab)
        return instance

    def _load_vocab(self, vocab_config: dict | str | None) -> dict:
        if isinstance(vocab_config, str):
            with open(vocab_config, encoding="utf-8") as fp:
                config = json.load(fp)
            return config["vocab"]
        if isinstance(vocab_config, dict):
            return vocab_config["vocab"]
        return {}

    def _create_audio(
        self,
        phonemes: str,
        voice: NDArray[np.float32],
        speed: float,
    ) -> tuple[NDArray[np.float32], int]:
        log.debug(f"Phonemes: {phonemes}")
        if len(phonemes) > MAX_PHONEME_LENGTH:
            log.warning(f"Phonemes are too long, truncating to {MAX_PHONEME_LENGTH} phonemes")
            phonemes = phonemes[:MAX_PHONEME_LENGTH]

        start_t = time.time()
        tokens = np.array(self.tokenizer.tokenize(phonemes), dtype=np.int64)
        assert len(tokens) <= MAX_PHONEME_LENGTH, (
            f"Context length is {MAX_PHONEME_LENGTH}, but leave room for the pad token 0 at the start & end"
        )
        voice = voice[len(tokens)]
        tokens = [[0, *tokens, 0]]
        if "input_ids" in [i.name for i in self.sess.get_inputs()]:
            inputs = {
                "input_ids": tokens,
                "style": np.array(voice, dtype=np.float32),
                "speed": np.array([speed], dtype=np.int32),
            }
        else:
            inputs = {
                "tokens": tokens,
                "style": voice,
                "speed": np.ones(1, dtype=np.float32) * speed,
            }
        audio = self.sess.run(None, inputs)[0]
        audio_duration = len(audio) / SAMPLE_RATE
        create_duration = time.time() - start_t
        rtf = create_duration / audio_duration if audio_duration else 0.0
        log.debug(
            f"Created audio in length of {audio_duration:.2f}s for {len(phonemes)} phonemes in {create_duration:.2f}s (RTF: {rtf:.2f})"
        )
        return audio, SAMPLE_RATE

    def get_voice_style(self, name: str) -> NDArray[np.float32]:
        return self.voices[name]

    def _split_phonemes(self, phonemes: str) -> list[str]:
        words = re.split(r"([.,!?;])", phonemes)
        batched_phonemes: list[str] = []
        current_batch = ""
        for part in words:
            part = part.strip()
            if part:
                if len(current_batch) + len(part) + 1 >= MAX_PHONEME_LENGTH:
                    if current_batch.strip():
                        batched_phonemes.append(current_batch.strip())
                    current_batch = part
                else:
                    if part in ".,!?;":
                        current_batch += part
                    else:
                        if current_batch:
                            current_batch += " "
                        current_batch += part
        if current_batch:
            batched_phonemes.append(current_batch.strip())
        return batched_phonemes

    def create(
        self,
        text: str,
        voice: str | NDArray[np.float32],
        speed: float = 1.0,
        lang: str = "en-us",
        is_phonemes: bool = False,
        trim: bool = True,
    ) -> tuple[NDArray[np.float32], int]:
        assert 0.5 <= speed <= 2.0, "Speed should be between 0.5 and 2.0"
        if isinstance(voice, str):
            assert voice in self.voices, f"Voice {voice} not found in available voices"
            voice = self.get_voice_style(voice)

        start_t = time.time()
        phonemes = text if is_phonemes else self.tokenizer.phonemize(text, lang)
        batched_phonemes = self._split_phonemes(phonemes)
        audio = []
        log.debug(f"Creating audio for {len(batched_phonemes)} batches for {len(phonemes)} phonemes")
        for phoneme_batch in batched_phonemes:
            audio_part, _ = self._create_audio(phoneme_batch, voice, speed)
            if trim:
                audio_part, _ = trim_audio(audio_part)
            audio.append(audio_part)
        audio = np.concatenate(audio) if audio else np.array([], dtype=np.float32)
        log.debug(f"Created audio in {time.time() - start_t:.2f}s")
        return audio, SAMPLE_RATE

    async def create_stream(
        self,
        text: str,
        voice: str | NDArray[np.float32],
        speed: float = 1.0,
        lang: str = "en-us",
        is_phonemes: bool = False,
        trim: bool = True,
    ) -> AsyncGenerator[tuple[NDArray[np.float32], int], None]:
        assert 0.5 <= speed <= 2.0, "Speed should be between 0.5 and 2.0"
        if isinstance(voice, str):
            assert voice in self.voices, f"Voice {voice} not found in available voices"
            voice = self.get_voice_style(voice)

        phonemes = text if is_phonemes else self.tokenizer.phonemize(text, lang)
        batched_phonemes = self._split_phonemes(phonemes)
        queue: asyncio.Queue[tuple[NDArray[np.float32], int] | None] = asyncio.Queue()

        async def process_batches():
            for i, phoneme_batch in enumerate(batched_phonemes):
                loop = asyncio.get_event_loop()
                audio_part, sample_rate = await loop.run_in_executor(
                    None, self._create_audio, phoneme_batch, voice, speed
                )
                if trim:
                    audio_part, _ = trim_audio(audio_part)
                log.debug(f"Processed chunk {i} of stream")
                await queue.put((audio_part, sample_rate))
            await queue.put(None)

        asyncio.create_task(process_batches())
        while True:
            chunk = await queue.get()
            if chunk is None:
                break
            yield chunk

    def get_voices(self) -> list[str]:
        return list(sorted(self.voices.keys()))
