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

This module follows the interface used by ``tktts_voicevox.py`` and
``tktts_qwen3.py``.  It calls Irodori-TTS directly through its official
``InferenceRuntime`` API; a Gradio UI or HTTP server does not need to be
running.  The loaded runtime is cached by Irodori-TTS and reused for all
subsequent utterances.

Text splitting and reading/pronunciation conversion are intentionally left to
the caller.  ``speak_dialogue`` only applies the replacement dictionary passed
by tktts, exactly like the other backends.
"""

from __future__ import annotations

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


missing = []
for lib, import_name in [
    ("torch", "torch"),
    ("Irodori-TTS", "irodori_tts"),
]:
    try:
        __import__(import_name)
    except ImportError:
        missing.append(lib)

if missing:
    raise ImportError(
        "Error: Missing libraries: "
        + ", ".join(missing)
        + "\n  Run this module in the Irodori-TTS uv environment."
        + "\n  install: uv sync --extra cu128"
    )

import torch
from irodori_tts.inference_runtime import (
    RuntimeKey,
    SamplingRequest,
    clear_cached_runtime,
    default_runtime_device,
    download_hf_checkpoint,
    get_cached_runtime,
    save_wav,
)

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


TTS_ENGINE_NAME = "irodori"
DEFAULT_IRODORI_MODEL = "Aratako/Irodori-TTS-v4.1-Small"
DEFAULT_IRODORI_VOICE = "default"
DEFAULT_CAPTION = "落ち着いた自然な声で、明瞭に読み上げる。"
DEFAULT_DEVICE = "auto"
DEFAULT_PRECISION = "auto"
DEFAULT_NUM_STEPS = 40
DEFAULT_CFG_SCALE_TEXT = 3.5
DEFAULT_CFG_SCALE_CAPTION = 3.0
DEFAULT_CFG_SCALE_SPEAKER = 5.0
DEFAULT_SEED = 0


_AVAILABLE_VOICES = [
    {
        "name": DEFAULT_IRODORI_VOICE,
        "description": (
            "Irodori-TTS voice. Use irodori_caption and/or reference WAV "
            "files to define voice identity and speaking style."
        ),
    }
]


def _resolve_device(device: str | None = DEFAULT_DEVICE) -> str:
    """Resolve ``auto`` to Irodori-TTS's preferred available device."""
    if device is None or str(device).strip().lower() == "auto":
        return default_runtime_device()
    value = str(device).strip().lower()
    if value == "cuda:0":
        return "cuda"
    return value


def _resolve_precision(device: str, precision: str | None) -> str:
    """Resolve ``auto`` to bf16 on CUDA/XPU and fp32 elsewhere."""
    value = "auto" if precision is None else str(precision).strip().lower()
    if value == "auto":
        return "bf16" if device in {"cuda", "xpu"} else "fp32"
    aliases = {
        "bfloat16": "bf16",
        "bf16": "bf16",
        "float32": "fp32",
        "fp32": "fp32",
    }
    if value not in aliases:
        raise ValueError(f"Unsupported Irodori-TTS precision: {precision}")
    resolved = aliases[value]
    if resolved == "bf16" and device not in {"cuda", "xpu"}:
        raise ValueError(f"bf16 is not supported for Irodori-TTS device={device}")
    return resolved


def _resolve_checkpoint(model_id: str) -> str:
    """Resolve a local checkpoint or download a Hugging Face checkpoint."""
    candidate = Path(str(model_id)).expanduser()
    if candidate.suffix.lower() in {".pt", ".safetensors"}:
        if not candidate.is_file():
            raise FileNotFoundError(f"Irodori-TTS checkpoint not found: {candidate}")
        return str(candidate.resolve())
    return str(download_hf_checkpoint(str(model_id)))


def load_model(
    model_id: str = DEFAULT_IRODORI_MODEL,
    device: str = DEFAULT_DEVICE,
    precision: str = DEFAULT_PRECISION,
    *,
    codec_device: str | None = None,
    codec_precision: str | None = None,
    force_reload: bool = False,
):
    """Load and cache an Irodori-TTS inference runtime."""
    resolved_device = _resolve_device(device)
    resolved_codec_device = _resolve_device(codec_device or resolved_device)
    resolved_precision = _resolve_precision(resolved_device, precision)
    resolved_codec_precision = _resolve_precision(
        resolved_codec_device,
        precision if codec_precision is None else codec_precision,
    )
    checkpoint = _resolve_checkpoint(model_id)

    if force_reload:
        clear_cached_runtime()

    key = RuntimeKey(
        checkpoint=checkpoint,
        model_device=resolved_device,
        codec_repo="Aratako/Semantic-DACVAE-Japanese-32dim",
        model_precision=resolved_precision,
        codec_device=resolved_codec_device,
        codec_precision=resolved_codec_precision,
        compile_model=False,
        compile_dynamic=False,
    )
    runtime, reloaded = get_cached_runtime(key)
    status = "loaded" if reloaded else "reused"
    print(
        "tktts_irodori.load_model(): "
        f"{status}, model={model_id}, device={resolved_device}, "
        f"precision={resolved_precision}"
    )
    return runtime


def unload_models() -> None:
    """Release the runtime cached by Irodori-TTS."""
    clear_cached_runtime()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def get_available_voices_info(model=None):
    """Return the virtual default voice used by Irodori-TTS.

    Irodori-TTS does not contain a fixed speaker list.  Voice identity is
    controlled by reference audio, while style can also be controlled by a
    caption.  ``model`` is accepted for tktts interface compatibility.
    """
    del model
    return [voice.copy() for voice in _AVAILABLE_VOICES]


def get_available_voices(model=None):
    """Return voice names compatible with the tktts voice selector."""
    return [voice["name"] for voice in get_available_voices_info(model)]


def list_available_voices(model=None) -> bool:
    """Print Irodori-TTS's virtual voice entry."""
    print(f"=== 利用可能な {TTS_ENGINE_NAME} voices ===")
    voices = get_available_voices_info(model)
    for voice in voices:
        print(f"  Name: {voice['name']}, Description: {voice['description']}")
    return bool(voices)


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


def resolve_speaker_id(speaker_name, voices_info=None, model=None):
    """Resolve ``default``/``irodori`` to the virtual Irodori voice."""
    del model
    if voices_info is None:
        voices_info = get_available_voices_info()

    query = _speaker_key(speaker_name)
    aliases = {
        "default",
        "irodori",
        "irodoritts",
        "標準",
        "デフォルト",
    }
    if query in aliases:
        return voices_info, DEFAULT_IRODORI_VOICE

    for voice in voices_info:
        if query == _speaker_key(voice["name"]):
            return voices_info, voice["name"]

    raise ValueError(
        "❌ Error in tktts_irodori.resolve_speaker_id(): "
        f"話者 [{speaker_name}] が見つかりませんでした。"
        " Irodori-TTSでは voice=default を使用し、"
        "irodori_captionまたは参照WAVで声を指定してください。"
    )


def _as_optional_int(value: Any, default: int | None) -> int | None:
    if value is None:
        return default
    if isinstance(value, str):
        text = value.strip().lower()
        if text in {"", "none", "random"}:
            return None
    return int(value)


def _as_ref_wavs(value: Any) -> list[str]:
    """Normalize a path, iterable, or semicolon-separated paths."""
    if value is None:
        return []
    if isinstance(value, (str, os.PathLike)):
        text = str(value).strip()
        if not text:
            return []
        values = [part.strip() for part in text.split(";") if part.strip()]
    elif isinstance(value, Iterable):
        values = [str(part).strip() for part in value if str(part).strip()]
    else:
        values = [str(value).strip()]

    paths = []
    for item in values:
        path = Path(item).expanduser()
        if not path.is_file():
            raise FileNotFoundError(f"Irodori-TTS reference WAV not found: {path}")
        paths.append(str(path.resolve()))
    return paths


def speak(
    outfile,
    text,
    voice=DEFAULT_IRODORI_VOICE,
    speak_rate=None,
    speak_pitch=None,
    *,
    caption: str | None = DEFAULT_CAPTION,
    ref_wav: str | os.PathLike | None = None,
    ref_wavs: Iterable[str | os.PathLike] | str | None = None,
    model=None,
    model_id: str = DEFAULT_IRODORI_MODEL,
    device: str = DEFAULT_DEVICE,
    precision: str = DEFAULT_PRECISION,
    codec_device: str | None = None,
    codec_precision: str | None = None,
    num_steps: int = DEFAULT_NUM_STEPS,
    cfg_scale_text: float = DEFAULT_CFG_SCALE_TEXT,
    cfg_scale_caption: float = DEFAULT_CFG_SCALE_CAPTION,
    cfg_scale_speaker: float = DEFAULT_CFG_SCALE_SPEAKER,
    duration_scale: float = 1.0,
    seed: int | str | None = DEFAULT_SEED,
    lora_adapter: str | None = None,
):
    """Generate one WAV file with Irodori-TTS v4.1-Small."""
    text = str(text).strip()
    if not text:
        print("❌ tktts_irodori.speak(): 読み上げテキストが空です")
        return None

    if _speaker_key(voice) not in {
        "default",
        "irodori",
        "irodoritts",
        "標準",
        "デフォルト",
    }:
        print(
            f"    ** Warning: Irodori-TTSには固定話者 [{voice}] がないため "
            "voice=defaultとして扱います"
        )
    if speak_rate is not None or speak_pitch is not None:
        print(
            "    ** Warning: Irodori-TTSでは speak_rate/speak_pitch を"
            "直接指定できないため無視します"
        )

    references = _as_ref_wavs(ref_wavs)
    if ref_wav is not None:
        references = _as_ref_wavs(ref_wav) + references

    if model is None:
        model = load_model(
            model_id=model_id,
            device=device,
            precision=precision,
            codec_device=codec_device,
            codec_precision=codec_precision,
        )

    try:
        result = model.synthesize(
            SamplingRequest(
                text=text,
                caption=str(caption).strip() if caption else None,
                ref_wav=None,
                ref_wavs=references or None,
                no_ref=not references,
                ref_normalize_db=-16.0,
                ref_ensure_max=True,
                num_candidates=1,
                decode_mode="sequential",
                seconds=None,
                duration_scale=float(duration_scale),
                max_ref_seconds=None,
                num_steps=int(num_steps),
                seed=_as_optional_int(seed, DEFAULT_SEED),
                cfg_guidance_mode="independent",
                cfg_scale_text=float(cfg_scale_text),
                cfg_scale_caption=float(cfg_scale_caption),
                cfg_scale_speaker=(float(cfg_scale_speaker) if references else 0.0),
                cfg_min_t=0.5,
                cfg_max_t=1.0,
                context_kv_cache=True,
                t_schedule_mode="linear",
                sway_coeff=-1.0,
                trim_tail=True,
                lora_adapter=lora_adapter,
            ),
            log_fn=lambda message: print(f"    {message}"),
        )
        output_path = save_wav(outfile, result.audio, result.sample_rate)
    except Exception as exc:
        print(f"❌ Irodori-TTS生成エラー: {exc}")
        return None

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

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


def _cfg_value(cfg, name: str, default: Any = None) -> Any:
    if cfg is None:
        return default
    return getattr(cfg, name, default)


def speak_dialogue(
    dialogue,
    replacements,
    target_voices,
    speakers=None,
    temp_dir=".",
    outfile=None,
    ext="wav",
    cfg=None,
    *,
    caption: str | None = None,
    ref_wav: str | os.PathLike | None = None,
    ref_wavs: Iterable[str | os.PathLike] | str | None = None,
    model=None,
    model_id: str = DEFAULT_IRODORI_MODEL,
    device: str = DEFAULT_DEVICE,
    precision: str = DEFAULT_PRECISION,
):
    """Generate temporary WAV files for a tktts dialogue sequence."""
    del outfile  # Kept in the signature for tktts backend compatibility.
    speakers = {} if speakers is None else speakers

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

    model_id = _cfg_value(cfg, "irodori_model_id", model_id)
    device = _cfg_value(cfg, "irodori_device", device)
    precision = _cfg_value(cfg, "irodori_precision", precision)
    codec_device = _cfg_value(cfg, "irodori_codec_device", None)
    codec_precision = _cfg_value(cfg, "irodori_codec_precision", None)

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

    effective_caption = (
        caption
        if caption is not None
        else _cfg_value(cfg, "irodori_caption", DEFAULT_CAPTION)
    )
    effective_ref_wav = (
        ref_wav
        if ref_wav is not None
        else _cfg_value(cfg, "irodori_ref_wav", None)
    )
    effective_ref_wavs = (
        ref_wavs
        if ref_wavs is not None
        else _cfg_value(cfg, "irodori_ref_wavs", None)
    )
    num_steps = _cfg_value(cfg, "irodori_num_steps", DEFAULT_NUM_STEPS)
    cfg_scale_text = _cfg_value(
        cfg, "irodori_cfg_scale_text", DEFAULT_CFG_SCALE_TEXT
    )
    cfg_scale_caption = _cfg_value(
        cfg, "irodori_cfg_scale_caption", DEFAULT_CFG_SCALE_CAPTION
    )
    cfg_scale_speaker = _cfg_value(
        cfg, "irodori_cfg_scale_speaker", DEFAULT_CFG_SCALE_SPEAKER
    )
    duration_scale = _cfg_value(cfg, "irodori_duration_scale", 1.0)
    seed = _cfg_value(cfg, "irodori_seed", DEFAULT_SEED)
    lora_adapter = _cfg_value(cfg, "irodori_lora_adapter", None)

    tmpfiles = []
    voices_info = get_available_voices_info(model)
    idx = 1
    is_monologue = bool(_cfg_value(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_IRODORI_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}")

            generated = speak(
                outfile=tmpfile,
                text=text,
                voice=target_voice,
                speak_rate=_cfg_value(cfg, "fspeak_rate", None),
                speak_pitch=_cfg_value(cfg, "fspeak_pitch", None),
                caption=effective_caption,
                ref_wav=effective_ref_wav,
                ref_wavs=effective_ref_wavs,
                model=model,
                num_steps=num_steps,
                cfg_scale_text=cfg_scale_text,
                cfg_scale_caption=cfg_scale_caption,
                cfg_scale_speaker=cfg_scale_speaker,
                duration_scale=duration_scale,
                seed=seed,
                lora_adapter=lora_adapter,
            )
            if generated is None:
                return False, tmpfiles

            tmpfiles.append(tmpfile)
            idx += 1

    return True, tmpfiles


if __name__ == "__main__":
    list_available_voices()