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

提供三種試用情境，並一律套用模型官方建議的最佳生成參數：
  1. 指定語者：李宏毅 —— 用模型內附、已取得授權的語者向量控制音色
  2. 聲音複製       —— 從參考音檔抽取語者向量（ECAPA centroid），走 speaker_centroid 路徑
  3. 長文逐句串流   —— 把長文切句，合成一句播一句，做出串流效果（每句含自動重試）

註：本 checkpoint 只訓練過 speaker_centroid（內附李宏毅向量同一機制）與 prompt 接續兩條
路徑；transcript-free 的 reference_wav_path 路徑從未餵過訓練資料（data._make_sample 丟掉
ref_audio），會吐出亂掉內容。故聲音複製改為「從錄音抽 centroid → speaker_centroid」——
免逐字稿、且走的是有訓練、可靠的路徑。

最佳參數讀自模型發佈中繼資料 release_metadata.json 的 recommended_generation_defaults
（以 500 句難句 + ASR CER 調出，最佳 trial：cfg=2.8 / steps=9），確保 demo 永遠跟著模型走。
"""

import json
import os
import time

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

# --------------------------------------------------------------------------- #
# 最佳生成參數：讀自模型發佈中繼資料；缺檔時退回官方建議值。
# --------------------------------------------------------------------------- #
_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.8,
    "inference_timesteps": 9,
    "max_len": 2000,
    "retry_badcase": True,
    "retry_badcase_max_times": 3,
    "retry_badcase_ratio_threshold": 6.0,
}
SPEAKER_ID = "hung_yi_lee"

try:
    with open(os.path.join(MODEL_DIR, "release_metadata.json"), encoding="utf-8") as f:
        _meta = json.load(f)
    _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 best params")
    GEN_KWARGS = dict(_BEST_FALLBACK)

print(f"[BlueMagpie] best generation params: {GEN_KWARGS}")

# 語者向量（已取得本人授權）
_centroids = torch.load(
    os.path.join(MODEL_DIR, "checkpoints", "hung_yi_lee_speaker_centroids.pt"),
    map_location="cpu",
    weights_only=True,
)
SPK_CENTROID = _centroids["centroids"][_centroids["speaker_ids"].index(SPEAKER_ID)]
print(f"[BlueMagpie] ready. sample_rate={SR}, speaker={SPEAKER_ID}")


# --------------------------------------------------------------------------- #
# 聲音複製：從參考音檔抽取語者向量（ECAPA），走「已訓練、可靠」的 speaker_centroid
# 路徑。與訓練用的 build_speaker_centroids.py 完全一致：同一個 ECAPA 模型、每段
# embedding 先 L2-normalize、取平均、再 L2-normalize（192 維），確保與內附李宏毅向量
# 同一向量空間。ECAPA 在 CPU 上跑、lazy 載入一次（不佔 GPU 時間配額）。
# --------------------------------------------------------------------------- #
_ECAPA = None


def _get_ecapa():
    global _ECAPA
    if _ECAPA is None:
        from speechbrain.inference.speaker import EncoderClassifier

        _ECAPA = EncoderClassifier.from_hparams(
            source="speechbrain/spkrec-ecapa-voxceleb", run_opts={"device": "cpu"}
        )
    return _ECAPA


@torch.no_grad()
def extract_centroid(wav_path: str) -> torch.Tensor:
    """參考音檔 -> [192] 語者向量（與訓練 centroid 同空間、同處理）。"""
    import librosa

    wav, _ = librosa.load(wav_path, sr=16000, mono=True)          # 16 kHz 單聲道
    x = torch.from_numpy(np.ascontiguousarray(wav)).float()
    if x.numel() < 16000:                                         # 太短（<1s）難抽穩定向量
        raise ValueError("參考音檔太短，請提供約 3–10 秒的乾淨語音。")
    clf = _get_ecapa()
    win = 6 * 16000                                               # 切約 6 秒一段，分段抽再平均
    chunks = [x[i : i + win] for i in range(0, x.numel(), win)]
    chunks = [c for c in chunks if c.numel() >= 16000] or [x]     # 每段至少 1 秒
    embs = []
    for c in chunks:
        e = clf.encode_batch(c.unsqueeze(0)).reshape(-1)          # [192]
        embs.append(torch.nn.functional.normalize(e, dim=0))
    mean = torch.stack(embs).mean(0)
    return torch.nn.functional.normalize(mean, dim=0).cpu()

# 註：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):
    text = (text or "").strip()
    if not text:
        raise gr.Error("請先輸入要合成的文字。")
    audio = model.generate(
        target_text=text,
        speaker_centroid=SPK_CENTROID,
        **GEN_KWARGS,
    )
    return (SR, _to_numpy(audio))


@gpu
def tts_clone(text, reference_audio):
    """聲音複製：從參考音檔抽出語者向量（ECAPA centroid），走 speaker_centroid 路徑。

    這條路徑在本 checkpoint 有被訓練（與內附李宏毅向量同一機制），可靠；**免逐字稿**。
    """
    text = (text or "").strip()
    if not text:
        raise gr.Error("請先輸入要合成的文字。")
    if not reference_audio:
        raise gr.Error("請先上傳或錄製一段參考音檔。")
    try:
        centroid = extract_centroid(reference_audio)
    except Exception as e:
        raise gr.Error(f"無法從參考音檔抽取語者向量：{e}")
    audio = model.generate(
        target_text=text,
        speaker_centroid=centroid,
        **GEN_KWARGS,
    )
    return (SR, _to_numpy(audio))


@gpu
def tts_stream(text):
    """長文逐句串流（李宏毅語者向量）：把長文切成句子，逐句合成、合成一句就播一句。

    逐句串流直到接近單次 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("找不到可合成的句子。")
    total = len(sentences)
    sentences = sentences[:_MAX_STREAM_SENTENCES]
    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=SPK_CENTROID, **GEN_KWARGS
        )
        yield (SR, _pcm16(_to_numpy(audio)))
        done += 1
    if done < total:
        gr.Warning(
            f"受單次 GPU 時間配額限制，本次串流前 {done} 句（共 {total} 句）；"
            "想念完整段可分批貼上。"
        )


# --------------------------------------------------------------------------- #
# 介面
# --------------------------------------------------------------------------- #
EXAMPLE_TEXTS = [
    "今天天氣真好，我們一起去散步吧。",
    "這是 AI TTS code switching 測試，混合中英文也沒問題。",
    "台灣藍鵲是一種叫聲響亮、辨識度很高的鳥。",
]

STREAM_EXAMPLES = [
    "台灣藍鵲是台灣特有的鳥類，羽色以亮藍為主，尾羽修長。牠們常成群在樹林間活動，"
    "叫聲響亮而容易辨認。今天天氣真好，我們一起出去走走，順便看看這些美麗的鳥吧！",
    "歡迎使用這個示範。這是一段比較長的文字，會被切成好幾句。"
    "系統會合成一句、播放一句，做出串流的效果。希望你會喜歡這個聲音！",
]

_PARAMS_NOTE = (
    f"已套用模型官方建議最佳參數："
    f"`cfg_value={GEN_KWARGS.get('cfg_value')}`、"
    f"`inference_timesteps={GEN_KWARGS.get('inference_timesteps')}`、"
    f"`retry_badcase={'on' if GEN_KWARGS.get('retry_badcase') else 'off'}`。"
)

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

**OpenFormosa Blue Magpie TTS** —— 支援**台灣華語**與**中英混合**的文字轉語音模型，輸出 48 kHz 單聲道語音。
模型：[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.Tab("指定語者：李宏毅"):
        gr.Markdown(
            "使用模型內附、**已取得李宏毅老師本人授權**的語者向量來控制音色"
            "（官方最佳參數正是針對此語者調校）。"
        )
        with gr.Row():
            with gr.Column():
                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_out)

    with gr.Tab("聲音複製"):
        gr.Markdown(
            "上傳或錄製一段參考音檔，系統會從中**抽取語者向量（音色指紋）**，再用它合成你的文字"
            "——**不需要逐字稿**。建議參考音檔約 3–10 秒、乾淨單一語者。\n\n"
            "（原理：和內附的李宏毅語者向量走同一條、**已訓練**的 `speaker_centroid` 路徑；"
            "用 ECAPA-TDNN 從你的音檔算出同樣的 192 維向量。）\n\n"
            "⚠️ **請只使用你已取得授權的聲音**，請勿在未經本人同意下複製他人聲音。"
        )
        with gr.Row():
            with gr.Column():
                c_text = gr.Textbox(
                    label="要合成的文字", lines=3,
                    placeholder="輸入中文或中英混合的文字…",
                )
                c_ref = gr.Audio(
                    label="參考音檔（上傳或錄製，約 3–10 秒）", type="filepath",
                    sources=["upload", "microphone"],
                )
                c_btn = gr.Button("複製音色合成", variant="primary")
            with gr.Column():
                c_out = gr.Audio(label="合成結果", type="numpy")
        c_btn.click(tts_clone, [c_text, c_ref], c_out)

    with gr.Tab("長文逐句串流"):
        gr.Markdown(
            "**長文逐句串流**：把長文切成句子，合成一句就播一句，做出串流（邊聽邊等）效果。"
            "第一句很快就能聽到，後面邊播邊合成。用**李宏毅語者向量**。\n\n"
            "> ZeroGPU 約 0.44x 實時，句與句之間可能有短暫間隔；每句都套用最佳參數（含自動重試）。"
            "單次串流以 ZeroGPU 時間配額為限（約可串數十句），超長會自動截斷並提示分批。"
        )
        with gr.Row():
            with gr.Column():
                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_out)

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


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