"""Qwen3-TTS backend for the tktts_* library family.

This module provides an interface compatible with ``tktts_voicevox.py`` while
running Qwen3-TTS locally through the official ``qwen-tts`` Python package.
The model is loaded lazily and cached, so consecutive calls do not reload it.
"""

from __future__ import annotations

import os
import re
from pathlib import Path
from typing import Any


missing = []
for lib, import_name in [
    ("torch", "torch"),
    ("soundfile", "soundfile"),
    ("qwen-tts", "qwen_tts"),
]:
    try:
        __import__(import_name)
    except ImportError:
        missing.append(lib)

if missing:
    raise ImportError(
        "Error: Missing libraries: "
        + ", ".join(missing)
        + "\n  install: pip install qwen-tts soundfile"
    )

import soundfile as sf
import torch
from qwen_tts import Qwen3TTSModel

try:
    from tktts_base import apply_replacements, normalize_speaker, split_dialogue
except ImportError as exc:
    raise ImportError(
        "tktts_qwen3 requires tktts_base.py to be importable."
    ) from exc


TTS_ENGINE_NAME = "qwen3"
DEFAULT_QWEN3_MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
DEFAULT_QWEN3_VOICE = "Ono_Anna"
DEFAULT_LANGUAGE = "Japanese"
DEFAULT_DEVICE = "auto"
DEFAULT_DTYPE = "auto"


# Official speakers included in the Qwen3-TTS CustomVoice checkpoints.
_AVAILABLE_VOICES = [
    {
        "name": "Vivian",
        "native_language": "Chinese",
        "description": "Bright, slightly edgy young female voice",
    },
    {
        "name": "Serena",
        "native_language": "Chinese",
        "description": "Warm, gentle young female voice",
    },
    {
        "name": "Uncle_Fu",
        "native_language": "Chinese",
        "description": "Seasoned male voice with a low, mellow timbre",
    },
    {
        "name": "Dylan",
        "native_language": "Chinese",
        "description": "Youthful Beijing male voice with a clear timbre",
    },
    {
        "name": "Eric",
        "native_language": "Chinese",
        "description": "Lively Chengdu male voice with husky brightness",
    },
    {
        "name": "Ryan",
        "native_language": "English",
        "description": "Dynamic male voice with strong rhythmic drive",
    },
    {
        "name": "Aiden",
        "native_language": "English",
        "description": "Sunny American male voice with a clear midrange",
    },
    {
        "name": "Ono_Anna",
        "native_language": "Japanese",
        "description": "Playful Japanese female voice with a light timbre",
    },
    {
        "name": "Sohee",
        "native_language": "Korean",
        "description": "Warm Korean female voice with rich emotion",
    },
]


_MODEL_CACHE: dict[tuple[str, str, str], Qwen3TTSModel] = {}


def _resolve_device(device: str | None = DEFAULT_DEVICE) -> str:
    """Resolve ``auto`` to an available CUDA device or CPU."""
    if device is None or str(device).strip().lower() == "auto":
        return "cuda:0" if torch.cuda.is_available() else "cpu"
    return str(device)


def _resolve_dtype(device: str, dtype: str | torch.dtype | None) -> torch.dtype:
    """Resolve a dtype name to a torch dtype suitable for the device."""
    if isinstance(dtype, torch.dtype):
        return dtype

    name = "auto" if dtype is None else str(dtype).strip().lower()
    if name == "auto":
        if device.startswith("cuda"):
            return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        return torch.float32

    dtype_map = {
        "bfloat16": torch.bfloat16,
        "bf16": torch.bfloat16,
        "float16": torch.float16,
        "fp16": torch.float16,
        "float32": torch.float32,
        "fp32": torch.float32,
    }
    if name not in dtype_map:
        raise ValueError(f"Unsupported dtype: {dtype}")
    return dtype_map[name]


def load_model(
    model_id: str = DEFAULT_QWEN3_MODEL,
    device: str = DEFAULT_DEVICE,
    dtype: str | torch.dtype = DEFAULT_DTYPE,
    *,
    force_reload: bool = False,
) -> Qwen3TTSModel:
    """Load and cache a Qwen3-TTS model.

    FlashAttention is intentionally not requested. The regular PyTorch
    attention path works on Windows without a separately installed CUDA SDK.
    """
    resolved_device = _resolve_device(device)
    resolved_dtype = _resolve_dtype(resolved_device, dtype)
    cache_key = (model_id, resolved_device, str(resolved_dtype))

    if force_reload:
        _MODEL_CACHE.pop(cache_key, None)

    if cache_key not in _MODEL_CACHE:
        print(
            "tktts_qwen3.load_model(): "
            f"model={model_id}, device={resolved_device}, dtype={resolved_dtype}"
        )
        _MODEL_CACHE[cache_key] = Qwen3TTSModel.from_pretrained(
            model_id,
            device_map=resolved_device,
            dtype=resolved_dtype,
        )

    return _MODEL_CACHE[cache_key]


def unload_models() -> None:
    """Release all models cached by this module."""
    _MODEL_CACHE.clear()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def get_available_voices_info(model: Qwen3TTSModel | None = None):
    """Return available speaker names and descriptions.

    When a model instance is supplied, its reported speaker list is used so
    that future CustomVoice checkpoints remain compatible.
    """
    if model is None:
        return [voice.copy() for voice in _AVAILABLE_VOICES]

    get_speakers = getattr(model, "get_supported_speakers", None)
    if not callable(get_speakers):
        return [voice.copy() for voice in _AVAILABLE_VOICES]

    known = {voice["name"].lower(): voice for voice in _AVAILABLE_VOICES}
    voices = []
    for speaker in get_speakers():
        info = known.get(str(speaker).lower(), {})
        voices.append(
            {
                "name": str(speaker),
                "native_language": info.get("native_language", ""),
                "description": info.get("description", ""),
            }
        )
    return voices


def get_available_voices(model: Qwen3TTSModel | None = None):
    """Return the available speaker names."""
    return [voice["name"] for voice in get_available_voices_info(model)]


def list_available_voices(model: Qwen3TTSModel | None = None) -> bool:
    """Print the available speakers."""
    print(f"=== 利用可能な {TTS_ENGINE_NAME} voices ===")
    voices = get_available_voices_info(model)
    if not voices:
        return False

    for voice in voices:
        language = voice.get("native_language") or "-"
        description = voice.get("description") or "-"
        print(
            f"  Name: {voice['name']}, "
            f"Native language: {language}, Description: {description}"
        )
    return True


def _speaker_key(name: str) -> str:
    normalized = normalize_speaker(str(name))
    return re.sub(r"[\s_\-（）()]+", "", normalized).lower()


def resolve_speaker_id(
    speaker_name: str,
    voices_info=None,
    model: Qwen3TTSModel | None = None,
):
    """Resolve a full or partial speaker name to a Qwen speaker name."""
    if voices_info is None:
        voices_info = get_available_voices_info(model)

    query = _speaker_key(speaker_name)
    if not query:
        raise ValueError("話者名が空です")

    # Prefer an exact normalized match before accepting a partial match.
    for voice in voices_info:
        if query == _speaker_key(voice["name"]):
            return voices_info, voice["name"]
    for voice in voices_info:
        if query in _speaker_key(voice["name"]):
            return voices_info, voice["name"]

    raise ValueError(
        "❌ Error in tktts_qwen3.resolve_speaker_id(): "
        f"話者 [{speaker_name}] が見つかりませんでした"
    )


def speak(
    outfile,
    text,
    voice=DEFAULT_QWEN3_VOICE,
    speak_rate=None,
    speak_pitch=None,
    *,
    language: str = DEFAULT_LANGUAGE,
    model: Qwen3TTSModel | None = None,
    model_id: str = DEFAULT_QWEN3_MODEL,
    device: str = DEFAULT_DEVICE,
    dtype: str | torch.dtype = DEFAULT_DTYPE,
    instruct: str | None = None,
):
    """Generate one audio file with a Qwen3-TTS CustomVoice model.

    ``speak_rate`` and ``speak_pitch`` are retained for interface compatibility
    with tktts_voicevox. CustomVoice does not expose equivalent parameters, so
    non-default values are currently ignored.
    """
    text = str(text).strip()
    if not text:
        print("❌ tktts_qwen3.speak(): 読み上げテキストが空です")
        return None

    if speak_rate is not None or speak_pitch is not None:
        print(
            "    ** Warning: Qwen3-TTS CustomVoiceでは "
            "speak_rate/speak_pitchを直接指定できないため無視します"
        )

    if model is None:
        model = load_model(model_id=model_id, device=device, dtype=dtype)

    kwargs: dict[str, Any] = {
        "text": text,
        "language": language,
        "speaker": voice,
    }
    if instruct:
        kwargs["instruct"] = instruct

    try:
        with torch.inference_mode():
            wavs, sample_rate = model.generate_custom_voice(**kwargs)
    except Exception as exc:
        print(f"❌ Qwen3-TTS生成エラー: {exc}")
        return None

    output_path = Path(outfile)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    sf.write(str(output_path), wavs[0], sample_rate)

    if output_path.exists():
        print(f"    ** 一時ファイル [{output_path}] を保存しました")
        return str(output_path)

    print(f"    ** Error: ファイル [{output_path}] の出力に失敗しました")
    return None


def speak_dialogue(
    dialogue,
    replacements,
    target_voices,
    speakers=None,
    temp_dir=".",
    outfile=None,
    ext="wav",
    cfg=None,
    *,
    language: str = DEFAULT_LANGUAGE,
    model: Qwen3TTSModel | None = None,
    model_id: str = DEFAULT_QWEN3_MODEL,
    device: str = DEFAULT_DEVICE,
    dtype: str | torch.dtype = DEFAULT_DTYPE,
):
    """Generate temporary audio files for a tktts dialogue sequence."""
    del outfile  # Kept in the signature for compatibility with tktts_voicevox.
    speakers = {} if speakers is None else speakers

    print("tktts_qwen3.speak_dialogue(): target_voices:", target_voices)

    if model is None:
        try:
            model = load_model(model_id=model_id, device=device, dtype=dtype)
        except Exception as exc:
            print(f"❌ Qwen3-TTSモデル読み込みエラー: {exc}")
            return False, []

    tmpfiles = []
    voices_info = get_available_voices_info(model)
    idx = 1
    is_monologue = bool(getattr(cfg, "monologue", False))

    for i, dialogue_item in enumerate(dialogue):
        print()
        print(f"Dialogue {i:04d}:")
        dialogue_list = split_dialogue(
            dialogue_item,
            target_voices,
            speakers=speakers,
            default_voice=DEFAULT_QWEN3_VOICE,
            is_monologue=is_monologue,
        )

        for speaker, text in dialogue_list:
            tmpfile = os.path.join(temp_dir, f"tmp_{idx:03d}.{ext}")
            text = apply_replacements(text, replacements)
            if isinstance(target_voices, str):
                speaker = target_voices

            try:
                voices_info, target_voice = resolve_speaker_id(
                    speaker,
                    voices_info=voices_info,
                    model=model,
                )
            except ValueError as exc:
                print(exc)
                return False, tmpfiles

            print(f"  {idx:04d}: voice={speaker} ({target_voice}): {text}")

            speak_rate = getattr(cfg, "fspeak_rate", None)
            speak_pitch = getattr(cfg, "fspeak_pitch", None)
            instruct = getattr(cfg, "qwen3_instruct", None)

            generated = speak(
                outfile=tmpfile,
                text=text,
                voice=target_voice,
                speak_rate=speak_rate,
                speak_pitch=speak_pitch,
                language=language,
                model=model,
                instruct=instruct,
            )
            if generated is None:
                return False, tmpfiles

            tmpfiles.append(tmpfile)
            idx += 1

    return True, tmpfiles


if __name__ == "__main__":
    list_available_voices()