"""Irodori-TTS Compare: hear Irodori-TTS v4.1 Small variants side by side (Hugging Face ZeroGPU Space).

Variants:
  Irodori-TTS-v4.1-Small                   upstream model (Aratako)
  Irodori-TTS-v4.1-Small-Yomi              kanji-reading fine-tune; served from the Yomi-Tech-tuned file with its
                                           reading frontend off (same tensors, byte-identical output)
  Irodori-TTS-v4.1-Small-Yomi-Tech-tuned   + technical-term reading frontend stored in the checkpoint

Local run: `python app.py` (set OURS_CHECKPOINT / BASE_CHECKPOINT to local model.safetensors files to skip
downloads; a tokenizer/ folder must sit beside each file; GRADIO_SERVER_PORT picks the port).
"""

from __future__ import annotations

import os
import re
import sys
from pathlib import Path

try:
    import spaces
except ImportError:  # running outside Hugging Face Spaces
    class spaces:  # noqa: N801 - mimics the module
        @staticmethod
        def GPU(*args, **kwargs):
            return args[0] if args and callable(args[0]) else (lambda fn: fn)

import gradio as gr
import numpy as np
import pandas as pd
import torch
from huggingface_hub import hf_hub_download

ROOT = Path(__file__).resolve().parent
sys.path[:0] = [str(ROOT), str(ROOT / "tech_reading")]

from irodori_tts.inference_runtime import InferenceRuntime, RuntimeKey, SamplingRequest  # noqa: E402
from irodori_tts.text_normalization import normalize_text  # noqa: E402
from letter_split import acronym_like  # noqa: E402
from reading_module import build_split, read_bundle  # noqa: E402
from span_expander import SPACE_MARK, find_spans  # noqa: E402

BASE_REPO, BASE_REVISION = "Aratako/Irodori-TTS-v4.1-Small", "2b28324dc263ed5e6638b3cf3dd94c82ead07b4b"
OURS_REPO = os.environ.get("OURS_REPO", "pentacoxian-dev/Irodori-TTS-v4.1-Small-Yomi-Tech-tuned")
OURS_REVISION = os.environ.get("OURS_REVISION") or None
DEVICE = "cuda" if (os.environ.get("SPACES_ZERO_GPU") or torch.cuda.is_available()) else "cpu"

SMALL, YOMI, TECH = "Irodori-TTS-v4.1-Small", "Irodori-TTS-v4.1-Small-Yomi", "Irodori-TTS-v4.1-Small-Yomi-Tech-tuned"
NONE = "(none)"
VARIANTS = {SMALL: ("base", False), YOMI: ("ours", False), TECH: ("ours", True)}  # name -> (weights, frontend)
NOTES = {SMALL: "upstream model", YOMI: "kanji-reading fine-tune", TECH: "Yomi + technical-term reading (ours)"}

YOMI_EXAMPLES = [
    ["我輩は屡々世界の人としての日本人の覚悟に関して述ぶるところがあった。", "屡々 → しばしば (the upstream model tends to say るるる)"],
    ["フランツは何と思ってか、そのまま踵を旋らして、自分の住んでいる村の方へ帰った。", "踵 → きびす (upstream: かかと)"],
    ["ひとりの男が、長い樋をつたって、だんだん下へおりてくるのです。", "樋 → とい (upstream: ひ)"],
    ["爾来我が中村屋は三十余年を通じて、一回たりともコンミッションに悩まされたことはない。", "爾来 → じらい (upstream: じゅらい)"],
    ["百舌鳥の声が喧しい程城内に交錯している。", "百舌鳥 → もず"],
]
TECH_EXAMPLES = [
    ["倉庫から届くMQTTの通信量が急に増えました。", "MQTT → エムキューティーティー (all four letters)"],
    ["受付でIDを確認してから、API keyを再発行しました。", "ID → アイディー, API key → エーピーアイキー"],
    ["来週のscheduleをcalendarで共有します。", "schedule → スケジュール, calendar → カレンダー"],
    ["GitHubのPull requestにLGTMをもらいました。", "GitHub → ギットハブ, Pull request → プルリクエスト, LGTM → エルジーティーエム"],
    ["Google MeetでHugging Faceのモデルについて話しましょう。", "Google Meet → グーグルミート, Hugging Face → ハギングフェイス"],
    ["violinの練習のあと、Kafkaの設定を見直しました。", "violin → バイオリン, Kafka → カフカ"],
]
MIXED_EXAMPLES = [
    ["百舌鳥の鳴き声を録音して、Hugging Faceに公開しました。", "百舌鳥 → もず, Hugging Face → ハギングフェイス"],
    ["踵を返して、Slackでscheduleを確認しました。", "踵 → きびす, Slack → スラック, schedule → スケジュール"],
    ["明日の会議は午後三時から始まります。", "Everyday Japanese: every model should read it correctly"],
]
INTRO = f"""\
# Irodori-TTS Compare

Pick two Irodori-TTS v4.1 Small variants and hear the same sentence from each.

| Variant | What it is |
| :--- | :--- |
| [{SMALL}](https://huggingface.co/{BASE_REPO}) | the upstream model |
| [{YOMI}](https://huggingface.co/takuma104/Irodori-TTS-v4.1-Small-Yomi) | text-encoder fine-tune for kanji readings |
| [{TECH}](https://huggingface.co/{OURS_REPO}) | the Yomi weights plus a reading frontend stored in the model file. English and technical terms are rewritten to their katakana reading before synthesis; the Japanese text is untouched. |

Both sides use the same text, seed, voice and settings.
"""
NOTE = ("Same seed on both sides. When the models or readings differ, timing differs too, so the voice can vary a "
        "little. Listen for the readings.")
READING = re.compile(r"^(.*?)\s+\((dictionary|reader|letters)\)$")


def checkpoint(env: str, repo: str, revision: str | None) -> Path:
    """model.safetensors with its tokenizer/ folder beside it (the unresolved Hub snapshot path)."""
    if os.environ.get(env):
        return Path(os.environ[env])
    for name in ("tokenizer/tokenizer.json", "tokenizer/tokenizer_config.json"):
        hf_hub_download(repo, name, revision=revision)
    return Path(hf_hub_download(repo, "model.safetensors", revision=revision))


def load(path: Path) -> InferenceRuntime:
    return InferenceRuntime.from_key(RuntimeKey(checkpoint=str(path), model_device=DEVICE, model_precision="fp32",
                                                codec_device=DEVICE, codec_precision="fp32"))


OURS_PATH = checkpoint("OURS_CHECKPOINT", OURS_REPO, OURS_REVISION)
RUNTIMES = {"base": load(checkpoint("BASE_CHECKPOINT", BASE_REPO, BASE_REVISION)), "ours": load(OURS_PATH)}
_bundle = read_bundle(OURS_PATH)
if _bundle is None:
    raise SystemExit(f"{OURS_PATH} has no technical-term reading frontend")
FRONTEND = build_split(_bundle, RUNTIMES["ours"].tokenizer, add_bos=RUNTIMES["ours"].tokenizer.add_bos)


def readings(text: str) -> pd.DataFrame:
    """How the Yomi-Tech-tuned frontend reads each English term in the text (runs on CPU)."""
    columns = ["English term", "Read as (Yomi-Tech-tuned)", "From"]
    text = (text or "").strip()
    if not text:
        return pd.DataFrame([], columns=columns)
    runtime, s = RUNTIMES["ours"], FRONTEND
    ids, mask = runtime.tokenizer.batch_encode([normalize_text(text).strip()], max_length=int(runtime.default_text_max_len))
    row, valid = ids[0].tolist(), int(mask[0].sum())
    new_ids, _, maps = s.retokenize(ids, mask)
    new_row, changed = new_ids[0].tolist(), dict(maps[0])
    rows = []
    for span in find_spans(row, s.classes, 1 if s.add_bos else 0, valid):
        term = s.span_text(row, span)
        if span not in changed:
            rows.append([term, term, "unchanged"])
            continue
        pieces = [s.strings[i].replace(SPACE_MARK, " ").strip() for i in new_row[slice(*changed[span])]]
        if re.search(r"[゠-ヿ]", "".join(pieces)):
            hit = term in s.lookup or (not acronym_like(term) and term.lower() in s.lookup)
            rows.append([term, "".join(pieces), "dictionary" if hit else "reader"])
        else:
            rows.append([term, " ".join(p for p in pieces if p), "letters"])
    return pd.DataFrame(rows, columns=columns)


def to_int16(audio: torch.Tensor) -> np.ndarray:
    wave = audio.detach().float().cpu().numpy().squeeze()
    return (np.clip(wave, -1.0, 1.0) * 32767).astype(np.int16)


@spaces.GPU(duration=90)
def generate(text: str, model_a: str, model_b: str, ref_audio: str | None, caption: str, seed: float, steps: float):
    text = (text or "").strip()
    if not text:
        raise gr.Error("Enter some text.")
    request = SamplingRequest(text=text, caption=(caption or "").strip() or None, ref_wavs=[ref_audio] if ref_audio else None,
                              no_ref=not ref_audio, num_steps=int(steps), seed=int(seed))
    outputs, timings = [None, None], []
    for slot, name in enumerate((model_a, model_b)):
        if not name or name == NONE:
            continue
        weights, tech = VARIANTS[name]
        runtime = RUNTIMES[weights]
        uninstall = FRONTEND.install(runtime.model) if tech else None
        try:
            result = runtime.synthesize(request)
        finally:
            if uninstall is not None:
                uninstall()
        outputs[slot] = (result.sample_rate, to_int16(result.audios[0]))
        timings.append(f"{'AB'[slot]}: {result.audios[0].shape[-1] / result.sample_rate:.1f} s audio")
    status = f"seed {int(seed)} · {'reference voice' if ref_audio else 'no reference'} · {int(steps)} steps · " + " · ".join(timings)
    return outputs[0], outputs[1], status


with gr.Blocks(title="Irodori-TTS Compare") as demo:
    gr.Markdown(INTRO)
    with gr.Row():
        model_a = gr.Dropdown(list(VARIANTS), value=SMALL, label="Model A")
        model_b = gr.Dropdown([*VARIANTS, NONE], value=TECH, label="Model B")
    with gr.Row():
        with gr.Column(scale=3):
            text = gr.Textbox(label="Text", value=TECH_EXAMPLES[0][0], lines=3)
            hint = gr.Textbox(label="What to listen for", value=TECH_EXAMPLES[0][1], interactive=False)
        with gr.Column(scale=2):
            table = gr.Dataframe(readings(TECH_EXAMPLES[0][0]), label="English terms in the text", interactive=False, wrap=True)
    button = gr.Button("Generate", variant="primary")
    with gr.Row():
        audio_a = gr.Audio(label=f"A · {SMALL}", interactive=False)
        audio_b = gr.Audio(label=f"B · {TECH}", interactive=False)
    status = gr.Markdown(NOTE)
    gr.Examples(YOMI_EXAMPLES, inputs=[text, hint], label="Kanji readings (Yomi)")
    gr.Examples(TECH_EXAMPLES, inputs=[text, hint], label="English and technical terms (Tech)")
    gr.Examples(MIXED_EXAMPLES, inputs=[text, hint], label="Both, and plain Japanese")
    with gr.Accordion("Settings", open=False):
        with gr.Row():
            seed = gr.Number(label="Seed", value=0, precision=0)
            steps = gr.Slider(label="Sampling steps", minimum=10, maximum=60, step=1, value=40)
        caption = gr.Textbox(label="Caption (optional voice description, e.g. 落ち着いた女性の声)", value="")
        ref = gr.Audio(label="Reference voice (optional)", type="filepath", sources=["upload", "microphone"])
    gr.Markdown(f"Model, code and evaluation: [{OURS_REPO}](https://huggingface.co/{OURS_REPO}).")
    text.input(lambda _: "", inputs=text, outputs=hint, show_progress="hidden")
    text.change(readings, inputs=text, outputs=table, show_progress="hidden")
    model_a.change(lambda n: gr.Audio(label=f"A · {n}", value=None), inputs=model_a, outputs=audio_a)
    model_b.change(lambda n: gr.Audio(label=f"B · {n}", value=None), inputs=model_b, outputs=audio_b)
    button.click(generate, inputs=[text, model_a, model_b, ref, caption, seed, steps], outputs=[audio_a, audio_b, status])

if __name__ == "__main__":
    demo.queue(max_size=20).launch()  # host/port from GRADIO_SERVER_NAME / GRADIO_SERVER_PORT (set by Spaces)