import asyncio
import logging
import re
from pathlib import Path
from typing import Optional

import numpy as np
import torch
from wyoming.audio import AudioChunk, AudioStart, AudioStop
from wyoming.error import Error
from wyoming.event import Event
from wyoming.info import Describe, Info
from wyoming.server import AsyncEventHandler
from wyoming.tts import Synthesize

from style_bert_vits2.nlp import bert_models
from style_bert_vits2.constants import Languages
from style_bert_vits2.tts_model import TTSModel

_LOGGER = logging.getLogger(__name__)

_VOICE_LOCK = asyncio.Lock()
_MODEL: Optional[TTSModel] = None
_BERT_INITIALIZED: set[Languages] = set()

_BERT_MODEL_NAMES = {
    Languages.JP: "ku-nlp/deberta-v2-large-japanese-char-wwm",
    Languages.EN: "microsoft/deberta-v3-large",
    Languages.ZH: "hfl/chinese-roberta-wwm-ext-large",
}

_HIRAGANA_KATAKANA = re.compile(r"[\u3040-\u309F\u30A0-\u30FF]")
_CJK = re.compile(r"[\u4E00-\u9FFF]")


def _optimize_gpu():
    if not torch.cuda.is_available():
        return

    torch.cuda.empty_cache = lambda: None
    _LOGGER.info("Disabled torch.cuda.empty_cache() to prevent re-allocation on every inference")

    if torch.version.hip is not None and hasattr(torch.backends, 'miopen'):
        torch.backends.miopen.benchmark = True
        _LOGGER.info("Enabled MIOpen benchmark mode to cache convolution solver solutions")


def _detect_language(text: str) -> Languages:
    if _HIRAGANA_KATAKANA.search(text):
        return Languages.JP
    if _CJK.search(text):
        return Languages.ZH
    return Languages.EN


def _load_bert_for_language(language: Languages, device: str, half: bool = False) -> None:
    if language in _BERT_INITIALIZED:
        return

    model_name = _BERT_MODEL_NAMES[language]
    _LOGGER.info("Loading BERT model for %s (%s)", language.name, model_name)
    bert_models.load_model(language, model_name)
    bert_models.load_tokenizer(language, model_name)
    bert = bert_models.__loaded_models.get(language)
    if bert is not None:
        bert.eval()
        bert.to(device).float()
        _LOGGER.info("BERT model for %s moved to %s", language.name, device)

    _BERT_INITIALIZED.add(language)


def _find_model_files(model_dir: Path):
    model_dir = model_dir.resolve()
    safetensors = list(model_dir.glob("*.safetensors"))
    config = model_dir / "config.json"
    style = model_dir / "style_vectors.npy"

    if safetensors and config.exists():
        return safetensors[0], config, style if style.exists() else None

    for subdir in sorted(model_dir.iterdir()):
        if not subdir.is_dir():
            continue
        safetensors = list(subdir.glob("*.safetensors"))
        config = subdir / "config.json"
        style = subdir / "style_vectors.npy"
        if safetensors and config.exists():
            return safetensors[0], config, style if style.exists() else None

    raise FileNotFoundError(
        f"No .safetensors files found in {model_dir} or its subdirectories"
    )


def _load_model(model_dir: Path, device: str, half: bool = False) -> TTSModel:
    model_path, config_path, style_path = _find_model_files(model_dir)

    _LOGGER.info("Creating TTSModel (model=%s, config=%s, device=%s)",
                 model_path.name, config_path.name, device)

    model = TTSModel(
        model_path=model_path,
        config_path=config_path,
        style_vec_path=style_path,
        device=device,
    )

    _LOGGER.info("Loading model weights...")
    model.load()

    net_g = getattr(model, "_TTSModel__net_g", None)
    if net_g is not None:
        net_g = net_g.float()
        setattr(model, "_TTSModel__net_g", net_g)

    _LOGGER.info("Model loaded successfully")
    return model


class GLaDOSEventHandler(AsyncEventHandler):
    def __init__(
        self,
        wyoming_info: Info,
        model_dir: Path,
        device: str,
        *args,
        half: bool = False,
        preload: bool = False,
        **kwargs,
    ) -> None:
        super().__init__(*args, **kwargs)
        self.wyoming_info_event = wyoming_info.event()
        self.model_dir = model_dir
        self.device = device
        self.half = half

        _optimize_gpu()

        if preload:
            _LOGGER.info("Pre-loading model at startup...")
            global _MODEL
            _MODEL = _load_model(model_dir, device, half)

    async def handle_event(self, event: Event) -> bool:
        if Describe.is_type(event.type):
            await self.write_event(self.wyoming_info_event)
            return True

        if not Synthesize.is_type(event.type):
            return True

        synthesize = Synthesize.from_event(event)
        return await self._handle_synthesize(synthesize)

    async def _handle_synthesize(self, synthesize: Synthesize) -> bool:
        global _MODEL

        text = synthesize.text.strip()
        if not text:
            return True

        language = _detect_language(text)

        speaker_id = 0
        style = "Neutral"
        if synthesize.voice is not None and synthesize.voice.speaker:
            try:
                speaker_id = int(synthesize.voice.speaker)
            except ValueError:
                pass

        _LOGGER.info("Synthesizing: text='%s' language=%s speaker=%s style=%s",
                      text[:80], language.name, speaker_id, style)

        try:
            async with _VOICE_LOCK:
                if _MODEL is None:
                    _LOGGER.info("Loading GLaDOS model from %s on %s",
                                 self.model_dir, self.device)
                    _MODEL = _load_model(self.model_dir, self.device, self.half)

                _load_bert_for_language(language, self.device, self.half)

                sr, audio = await asyncio.to_thread(
                    _MODEL.infer,
                    text=text,
                    language=language,
                    speaker_id=speaker_id,
                    style=style,
                )

            audio_int16 = np.round(audio).astype(np.int16)
            raw_bytes = audio_int16.tobytes()

            rate = sr
            width = 2
            channels = 1

            await self.write_event(
                AudioStart(rate=rate, width=width, channels=channels).event()
            )

            samples_per_chunk = 1024
            bytes_per_sample = width * channels
            bytes_per_chunk = bytes_per_sample * samples_per_chunk

            for i in range(0, len(raw_bytes), bytes_per_chunk):
                chunk = raw_bytes[i:i + bytes_per_chunk]
                await self.write_event(
                    AudioChunk(
                        audio=chunk,
                        rate=rate,
                        width=width,
                        channels=channels,
                    ).event()
                )

            await self.write_event(AudioStop().event())
            return True

        except Exception as err:
            _LOGGER.exception("Synthesis failed")
            await self.write_event(
                Error(text=str(err), code=err.__class__.__name__).event()
            )
            return True