"""BlueMagpie-TTS 線上試用 Space (Gradio + ZeroGPU)。

提供三種試用情境，並一律套用模型官方建議的生成參數：
  1. 指定語者       —— 用模型內附、已取得授權的語者向量控制音色（李宏毅／女聲）
  2. 聲音複製       —— 一段 3 秒以上的參考音檔直接走 reference_wav_path，免逐字稿
                       （checkpoint step_0006000 起正式支援，內部評測 CER 8.99%）
  3. 長文逐句串流   —— 把長文切句，合成一句播一句，做出串流效果（每句含自動重試）

生成參數：優先讀模型發佈中繼資料 release_metadata.json 的
recommended_generation_defaults（若有）；新版中繼資料未提供時，退回模型卡建議值
（cfg=2.0 / steps=10 / retry_badcase=True），確保 demo 永遠跟著模型走。
"""

import json
import os
import time

import librosa
import numpy as np
import torch
import gradio as gr
from huggingface_hub import snapshot_download
from transformers import PreTrainedTokenizerFast

from bluemagpie import BlueMagpieModel

# --------------------------------------------------------------------------- #
# ZeroGPU：用 @spaces.GPU 動態取得 GPU；本機無 spaces 時退回 cuda/cpu 自動偵測。
# --------------------------------------------------------------------------- #
try:
    import spaces

    gpu = spaces.GPU(duration=120)
    DEVICE = "cuda"
except ImportError:
    def gpu(fn):
        return fn

    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

REPO_ID = "OpenFormosa/BlueMagpie-TTS"

# --------------------------------------------------------------------------- #
# 載入模型（啟動時做一次）
# --------------------------------------------------------------------------- #
print(f"[BlueMagpie] downloading model from {REPO_ID} ...")
MODEL_DIR = snapshot_download(REPO_ID)

tokenizer = PreTrainedTokenizerFast(
    tokenizer_file=os.path.join(MODEL_DIR, "tokenizer.json")
)
print(f"[BlueMagpie] loading model on device={DEVICE} ...")
model = BlueMagpieModel.from_local(
    MODEL_DIR, tokenizer=tokenizer, training=False, device=DEVICE
)
SR = model.sample_rate

# --------------------------------------------------------------------------- #
# 生成參數：優先讀模型發佈中繼資料（舊版才有 recommended_generation_defaults）；
# 新版（step_0006000 起）未提供時，退回模型卡建議值。
# --------------------------------------------------------------------------- #
_GEN_KEYS = {
    "cfg_value",
    "inference_timesteps",
    "max_len",
    "min_len",
    "retry_badcase",
    "retry_badcase_max_times",
    "retry_badcase_ratio_threshold",
}
_BEST_FALLBACK = {
    "cfg_value": 2.0,
    "inference_timesteps": 10,
    "max_len": 2000,
    "retry_badcase": True,
    "retry_badcase_max_times": 3,
    "retry_badcase_ratio_threshold": 6.0,
}
SPEAKER_ID = "hung_yi_lee"
CHECKPOINT = "?"

try:
    with open(os.path.join(MODEL_DIR, "release_metadata.json"), encoding="utf-8") as f:
        _meta = json.load(f)
    CHECKPOINT = _meta.get("checkpoint", CHECKPOINT)
    _rec = _meta.get("recommended_generation_defaults", {}) or {}
    GEN_KWARGS = {k: v for k, v in _rec.items() if k in _GEN_KEYS} or dict(_BEST_FALLBACK)
    SPEAKER_ID = _rec.get("speaker_id", SPEAKER_ID)
except Exception as e:  # 任何讀取/解析問題都退回安全的建議值
    print(f"[BlueMagpie] release_metadata.json unavailable ({e}); using fallback params")
    GEN_KWARGS = dict(_BEST_FALLBACK)

print(f"[BlueMagpie] checkpoint={CHECKPOINT}, generation params: {GEN_KWARGS}")

# 介面上可調整的兩個生成參數；預設即官方建議值。
DEFAULT_CFG = float(GEN_KWARGS.get("cfg_value", 2.0))
DEFAULT_STEPS = int(GEN_KWARGS.get("inference_timesteps", 10))


def _gen_kwargs(cfg_value=None, inference_timesteps=None):
    """以官方建議參數為底，套用使用者在介面上調整的 cfg / 取樣步數。"""
    kw = dict(GEN_KWARGS)
    if cfg_value is not None:
        kw["cfg_value"] = float(cfg_value)
    if inference_timesteps is not None:
        kw["inference_timesteps"] = int(inference_timesteps)
    return kw


def _apply_speed(wav, speed):
    """語速調整：保持音高的時間伸縮（librosa time_stretch，CPU 後處理）。

    模型本身沒有語速參數（語速由模型隱含決定，約 4 字/秒），故以後處理實作；
    1.0 為原速，直接回傳原波形。
    """
    speed = float(speed or 1.0)
    if abs(speed - 1.0) < 1e-3:
        return wav
    return librosa.effects.time_stretch(wav, rate=speed)

# 語者向量：載入多語者表（優先 speaker_centroids.pt），建成「顯示名稱 -> 向量」。
# 內附的李宏毅向量已取得本人授權；其餘為通用語者向量。
_DISPLAY = {"hung_yi_lee": "李宏毅", "female_voice": "女聲"}


def _load_speakers():
    for fn in ("speaker_centroids.pt", "hung_yi_lee_speaker_centroids.pt"):
        p = os.path.join(MODEL_DIR, "checkpoints", fn)
        if os.path.exists(p):
            blob = torch.load(p, map_location="cpu", weights_only=True)
            return {_DISPLAY.get(sid, sid): c for sid, c in zip(blob["speaker_ids"], blob["centroids"])}
    raise RuntimeError("no speaker centroid file found")


SPEAKERS = _load_speakers()
DEFAULT_SPEAKER = _DISPLAY.get(SPEAKER_ID, SPEAKER_ID)
if DEFAULT_SPEAKER not in SPEAKERS:
    DEFAULT_SPEAKER = next(iter(SPEAKERS))
SPK_CENTROID = SPEAKERS[DEFAULT_SPEAKER]  # 預設向量（長文串流沿用）
print(f"[BlueMagpie] ready. sample_rate={SR}, speakers={list(SPEAKERS)}")


def _audio_seconds(path):
    """回傳音檔長度（秒）；讀不到時回 0。"""
    import librosa

    try:
        return float(librosa.get_duration(path=path))
    except Exception:
        return 0.0

# 註：repo 的 bluemagpie.serving 加速引擎（torch.compile / 連續批次）是為「常駐
# 專用 GPU、追求吞吐」設計。在 ZeroGPU（GPU 逐請求序列化掛載）上實測後，單一
# 請求延遲沒有改善（warm ~7.0s vs eager ~6.2s），反而帶來首呼叫 ~50s 編譯成本與
# CUDA graph 反覆重編譯的尖峰，故此 demo 維持 eager。


def _to_numpy(audio: torch.Tensor):
    return audio.squeeze().float().cpu().numpy()


def _pcm16(x: np.ndarray) -> np.ndarray:
    """float 波形 -> 16-bit PCM，供 Gradio 串流輸出（格式最穩定）。"""
    return (np.clip(x, -1.0, 1.0) * 32767.0).astype(np.int16)


# 長文逐句串流：以句末標點／換行切句。逐句串流直到接近 ZeroGPU 單次 GPU 時間配額
# 為止（不是固定句數），盡量把整段長文串完；另設一個防濫用的硬上限。
_SENT_ENDERS = "。！？；!?…\n"
_GPU_BUDGET_S = 100         # @spaces.GPU(duration=120)，留約 20s 緩衝給最後一句
_MAX_STREAM_SENTENCES = 80  # 防濫用硬上限（通常時間配額會先到）


def _split_sentences(text):
    out, cur = [], []
    for ch in text:
        cur.append(ch)
        if ch in _SENT_ENDERS:
            s = "".join(cur).strip()
            if s:
                out.append(s)
            cur = []
    s = "".join(cur).strip()
    if s:
        out.append(s)
    # 把落單的標點（如連續「！？」被拆出的單一字元）併回前一句
    merged = []
    for s in out:
        if merged and len(s) <= 1:
            merged[-1] += s
        else:
            merged.append(s)
    return merged


# --------------------------------------------------------------------------- #
# 推論：三種模式都套用 GEN_KWARGS（官方建議參數）
# --------------------------------------------------------------------------- #
@gpu
def tts_speaker(text, speaker=DEFAULT_SPEAKER, cfg=DEFAULT_CFG, steps=DEFAULT_STEPS, speed=1.0):
    text = (text or "").strip()
    if not text:
        raise gr.Error("請先輸入要合成的文字。")
    audio = model.generate(
        target_text=text,
        speaker_centroid=SPEAKERS.get(speaker, SPK_CENTROID),
        **_gen_kwargs(cfg, steps),
    )
    return (SR, _apply_speed(_to_numpy(audio), speed))


@gpu
def _clone_gpu(text, ref_path, cfg, steps):
    audio = model.generate(
        target_text=text, reference_wav_path=ref_path, **_gen_kwargs(cfg, steps)
    )
    return _to_numpy(audio)


def tts_clone(text, recording, cfg=DEFAULT_CFG, steps=DEFAULT_STEPS, speed=1.0):
    """聲音複製（免逐字稿）：把參考音檔直接交給模型（reference_wav_path 路徑，
    checkpoint step_0006000 起正式支援）。一段 3 秒以上的乾淨語音即可。"""
    text = (text or "").strip()
    if not text:
        raise gr.Error("請先輸入要合成的文字。")
    if not recording:
        raise gr.Error("請先錄音或上傳參考音檔。")
    dur = _audio_seconds(recording)
    if 0 < dur < 3:
        gr.Warning(f"參考音檔只有約 {dur:.0f} 秒，建議 3 秒以上會更穩定。")
    wav = _clone_gpu(text, recording, cfg, steps)
    return (SR, _apply_speed(wav, speed))   # 變速在 GPU 函式外做，不佔 GPU 配額


@gpu
def tts_stream(text, speaker=DEFAULT_SPEAKER, cfg=DEFAULT_CFG, steps=DEFAULT_STEPS, speed=1.0):
    """長文逐句串流（指定語者向量）：把長文切成句子，逐句合成、合成一句就播一句。

    逐句串流直到接近單次 GPU 時間配額（約 100s）為止，盡量把整段長文串完，而非固定
    句數；真的超過才截斷並提示分批。每句都用一般 generate（含 retry_badcase）產生
    完整音檔再 yield；句與句之間可能有短暫間隔（ZeroGPU 約 0.44x 即時所致）。
    """
    text = (text or "").strip()
    if not text:
        raise gr.Error("請先輸入要合成的文字。")
    sentences = _split_sentences(text)
    if not sentences:
        raise gr.Error("找不到可合成的句子。")
    centroid = SPEAKERS.get(speaker, SPK_CENTROID)
    total = len(sentences)
    sentences = sentences[:_MAX_STREAM_SENTENCES]
    gen_kw = _gen_kwargs(cfg, steps)
    start = time.time()
    done = 0
    for sent in sentences:
        # 至少先產出第一句；之後一旦逼近 GPU 時間配額就停，避免配額被回收而報錯
        if done and time.time() - start > _GPU_BUDGET_S:
            break
        audio = model.generate(
            target_text=sent, speaker_centroid=centroid, **gen_kw
        )
        yield (SR, _pcm16(_apply_speed(_to_numpy(audio), speed)))
        done += 1
    if done < total:
        gr.Warning(
            f"受單次 GPU 時間配額限制，本次串流前 {done} 句（共 {total} 句）；"
            "想念完整段可分批貼上。"
        )


# --------------------------------------------------------------------------- #
# 介面
# --------------------------------------------------------------------------- #
# 第 2–5 句刻意使用台灣審定讀音與中國讀音不同的字詞（垃圾ㄌㄜˋㄙㄜˋ、
# 期ㄑㄧˊ、績ㄐㄧ、括ㄍㄨㄚ、究ㄐㄧㄡˋ、法國ㄈㄚˋ、企ㄑㄧˋ、髮ㄈㄚˇ、
# 說服ㄕㄨㄟˋ、骰子ㄊㄡˊ），可測試模型是否唸出台灣的讀法。
EXAMPLE_TEXTS = [
    "今天天氣真好，我們一起去散步吧。",
    "我要吃蚵仔煎然後去丟垃圾。",
    "這學期的成績包括研究報告和期末考。",
    "他星期五要去一家法國企業面試，記得把頭髮整理好。",
    "我說服朋友擲骰子決定誰洗碗，結果他輸了。",
    "這是 AI TTS code switching 測試，混合中英文也沒問題。",
    "台灣藍鵲是一種叫聲響亮、辨識度很高的鳥。",
]

STREAM_EXAMPLES = [
    "台灣藍鵲是台灣特有的鳥類，羽色以亮藍為主，尾羽修長。牠們常成群在樹林間活動，"
    "叫聲響亮而容易辨認。今天天氣真好，我們一起出去走走，順便看看這些美麗的鳥吧！",
    "放假的時候，我最喜歡去逛夜市。先來一份蚵仔煎，再買一杯珍珠奶茶，邊走邊喝。"
    "逛累了就坐捷運回家，路上還可以看看街頭藝人表演。台灣的夜市文化，真的會讓人一去再去！",
    "歡迎使用這個示範。這是一段比較長的文字，會被切成好幾句。"
    "系統會合成一句、播放一句，做出串流的效果。希望你會喜歡這個聲音！",
]

_PARAMS_NOTE = (
    f"模型 checkpoint：`{CHECKPOINT}`，預設套用官方建議參數："
    f"`cfg_value={DEFAULT_CFG}`、"
    f"`inference_timesteps={DEFAULT_STEPS}`、"
    f"`retry_badcase={'on' if GEN_KWARGS.get('retry_badcase') else 'off'}`；"
    f"可展開下方「⚙️ 生成參數」自行調整。"
)

HEADER = f"""
# 🐦‍⬛ BlueMagpie-TTS · 文字轉語音線上試用

**OpenFormosa Blue Magpie TTS** —— 針對**台灣華語**與**中英混合**打造的文字轉語音模型，輸出 48 kHz 語音；
支援**免逐字稿聲音複製**（內部評測：語者向量 CER 7.44%、參考音檔 CER 8.99%，較基準 11.45% 顯著降低）。
模型：[OpenFormosa/BlueMagpie-TTS](https://huggingface.co/OpenFormosa/BlueMagpie-TTS) ·
程式碼：[GitHub](https://github.com/OpenFormosa/BlueMagpie-TTS)

> {_PARAMS_NOTE}
"""

with gr.Blocks(title="BlueMagpie-TTS Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown(HEADER)

    with gr.Accordion("⚙️ 生成參數（預設為官方建議值，三個分頁共用）", open=False):
        with gr.Row():
            g_cfg = gr.Slider(
                1.0, 4.0, value=DEFAULT_CFG, step=0.1,
                label="cfg_value（引導強度）",
                info="越大越貼合語者條件、但可能較不自然（建議 2.0）",
            )
            g_steps = gr.Slider(
                4, 20, value=DEFAULT_STEPS, step=1,
                label="inference_timesteps（取樣步數）",
                info="越多通常品質越好、速度越慢（建議 10）",
            )
            g_speed = gr.Slider(
                0.8, 1.3, value=1.0, step=0.05,
                label="語速（倍率）",
                info="1.0 為原速；保持音高的變速後處理",
            )

    with gr.Tab("指定語者"):
        gr.Markdown(
            "用模型內附的**語者向量**控制音色。**李宏毅**老師的向量已取得本人授權；"
            "**女聲**為通用語者向量。"
        )
        with gr.Row():
            with gr.Column():
                s_spk = gr.Dropdown(
                    list(SPEAKERS), value=DEFAULT_SPEAKER, label="語者",
                )
                s_text = gr.Textbox(
                    label="要合成的文字", lines=3,
                    placeholder="輸入中文或中英混合的文字…",
                )
                s_btn = gr.Button("合成", variant="primary")
            with gr.Column():
                s_out = gr.Audio(label="合成結果", type="numpy")
        gr.Examples(EXAMPLE_TEXTS, inputs=s_text, label="範例文字")
        s_btn.click(tts_speaker, [s_text, s_spk, g_cfg, g_steps, g_speed], s_out)

    with gr.Tab("聲音複製"):
        gr.Markdown(
            "**聲音複製，免逐字稿**：錄或上傳一段 **3 秒以上**的乾淨語音，"
            "直接交給模型模仿該語者的音色，**不需要**參考音檔的文字稿。\n\n"
            "⚠️ **請只使用你已取得授權的聲音**，請勿在未經本人同意下複製他人聲音。"
        )
        with gr.Row():
            with gr.Column():
                c_text = gr.Textbox(
                    label="要合成的文字", lines=3,
                    placeholder="輸入中文或中英混合的文字…",
                )
                c_rec = gr.Audio(
                    label="錄音或上傳參考語音（3 秒以上）",
                    type="filepath", sources=["microphone", "upload"],
                )
                c_btn = gr.Button("複製音色合成", variant="primary")
            with gr.Column():
                c_out = gr.Audio(label="合成結果", type="numpy")
        c_btn.click(tts_clone, [c_text, c_rec, g_cfg, g_steps, g_speed], c_out)

    with gr.Tab("長文逐句串流"):
        gr.Markdown(
            "**長文逐句串流**：把長文切成句子，合成一句就播一句，做出串流（邊聽邊等）效果。"
            "第一句很快就能聽到，後面邊播邊合成。用**所選語者向量**。\n\n"
            "> ZeroGPU 約 0.44x 即時，句與句之間可能有短暫間隔；每句都套用所選的生成參數（含自動重試）。"
            "單次串流以 ZeroGPU 時間配額為限（約可串數十句），超長會自動截斷並提示分批。"
        )
        with gr.Row():
            with gr.Column():
                t_spk = gr.Dropdown(
                    list(SPEAKERS), value=DEFAULT_SPEAKER, label="語者",
                )
                t_text = gr.Textbox(
                    label="要合成的長文", lines=5,
                    placeholder="貼上一段多句的文字，會逐句合成、逐句播放…",
                )
                t_btn = gr.Button("逐句串流合成", variant="primary")
            with gr.Column():
                t_out = gr.Audio(label="逐句串流播放", streaming=True, autoplay=True)
        gr.Examples(STREAM_EXAMPLES, inputs=t_text, label="長文範例")
        t_btn.click(tts_stream, [t_text, t_spk, g_cfg, g_steps, g_speed], t_out)

    gr.Markdown(
        "---\n"
        "合成的語音僅供研究與評估展示用途，輸出可能不完美；正式使用前請人工檢視。"
        " · Apache-2.0"
    )


if __name__ == "__main__":
    demo.queue().launch()