"""本地 TTS（sherpa-onnx VITS 中文）。

选型依据（P0-语音基准报告.md §六之二）：

    vits-zh-ll        服务器 RTF 0.09 / N5095 RTF 1.86   ← 用户试听选定，音质最好
    piper-huayan-med  服务器 0.36 / N5095 0.32          （音质次之）
    kokoro            1.96 以上，排除

**两个必须写进代码的实测结论：**

1. **VITS 是非自回归的，没有"边合成边播"。**
   实测回调只在合成结束后触发一次，所以首音延迟 == 整句合成耗时。
   提前出声的唯一手段是**按语义短句逐句合成**——把长回复切成短句，
   第一句合成完就能播。

2. **`max_num_sentences=1` 会静默截断多句输入。**
   传含多个句号的文本只会合成第一句，**不报错**。
   第一版测"长段"时只出 0.94 秒音频，差点当成正常结果。
   所以这里必须先分句，逐句送。

另外按方案 §5 要求，输出是 16kHz 单声道 float32；调用方负责重采样。
"""

from __future__ import annotations

import logging
import re
from pathlib import Path
from typing import Iterator

import numpy as np
import sherpa_onnx

from ..config import VoiceConfig

log = logging.getLogger(__name__)

#: 句末标点。中英文都要，方案 §5 要求覆盖中英混读。
_SENT_END = "。！？；!?;"
#: 从句标点，长句再切一刀，避免单句合成太久
_CLAUSE = "，,、：:"

#: 单句最大字符数。超过就按从句标点再切——
#: 句子越长，首音延迟越大（首音 == 整句合成耗时）。
MAX_CHARS_PER_SEGMENT = 30


def text_for_speech(text: str) -> str:
    """去掉只适合屏幕阅读的链接/来源标记，避免 TTS 朗读网址。"""
    text = text or ""
    # 搜索回答约定把引用放在独立的“来源：”行；语音正文到此结束。
    text = re.split(
        r"(?:^|\n)\s*(?:参考来源|资料来源|来源)\s*[：:]",
        text,
        maxsplit=1,
        flags=re.IGNORECASE,
    )[0]
    # 即使模型偶尔用了 Markdown 链接，也只朗读标题。
    text = re.sub(r"\[([^\]]+)\]\(https?://[^)]+\)", r"\1", text)
    text = re.sub(r"https?://\S+", "", text)
    text = re.sub(r"[*#`|]", "", text)
    return " ".join(text.split()).strip()


class TTSError(RuntimeError):
    pass


def split_sentences(text: str, *, max_chars: int = MAX_CHARS_PER_SEGMENT) -> list[str]:
    """按语义切分，返回可逐句合成的片段。

    先按句末标点切，过长的再按从句标点切，仍过长就硬切。
    标点保留在片段末尾（TTS 需要它来判断语调）。

    这个函数存在的原因是 `max_num_sentences=1` 会静默丢内容——
    不自己切，就会"念到一半没了"。
    """
    text = (text or "").strip()
    if not text:
        return []

    parts = re.split(rf"(?<=[{_SENT_END}])", text)
    out: list[str] = []
    for p in parts:
        p = p.strip()
        if not p:
            continue
        if len(p) <= max_chars:
            out.append(p)
            continue
        # 过长：按从句标点再切
        subs = re.split(rf"(?<=[{_CLAUSE}])", p)
        buf = ""
        for s in subs:
            if not s:
                continue
            if len(buf) + len(s) <= max_chars:
                buf += s
            else:
                if buf:
                    out.append(buf)
                # 单从句仍然过长，硬切
                while len(s) > max_chars:
                    out.append(s[:max_chars])
                    s = s[max_chars:]
                buf = s
        if buf:
            out.append(buf)
    return out


class LocalTTS:
    def __init__(self, cfg: VoiceConfig):
        self.cfg = cfg
        d = Path(cfg.model_dir) / "tts"
        root = Path(cfg.model_dir)
        files = {
            "model": d / "model.onnx",
            "lexicon": d / "lexicon.txt",
            "tokens": d / "tokens.txt",
        }
        missing = [k for k, p in files.items() if not p.is_file()]
        if missing:
            raise TTSError(
                f"TTS 模型文件缺失：{missing}（目录 {d}）。"
                f"先跑 p0/stt_tts_bench/download_models.sh")

        rule_fsts = ",".join(
            str(d / f) for f in ("date.fst", "number.fst", "phone.fst", "new_heteronym.fst")
            if (d / f).is_file())

        self._tts = sherpa_onnx.OfflineTts(sherpa_onnx.OfflineTtsConfig(
            model=sherpa_onnx.OfflineTtsModelConfig(
                vits=sherpa_onnx.OfflineTtsVitsModelConfig(
                    model=str(files["model"]),
                    lexicon=str(files["lexicon"]),
                    tokens=str(files["tokens"]),
                    dict_dir=str(d / "dict"),
                ),
                num_threads=cfg.tts_threads,
                provider="cpu",
            ),
            rule_fsts=rule_fsts,
            # 逐句生成：这是唯一能提前出声的手段（见模块文档）
            max_num_sentences=1,
        ))
        self.sample_rate = self._tts.sample_rate
        log.info("TTS 就绪：%d 线程，输出 %d Hz", cfg.tts_threads, self.sample_rate)

    def synthesize(self, text: str) -> np.ndarray:
        """合成一段文本（内部已确保是单句）。"""
        if not text.strip():
            return np.zeros(0, dtype=np.float32)
        audio = self._tts.generate(text, sid=0, speed=1.0)
        return np.asarray(audio.samples, dtype=np.float32)

    def synthesize_sentences(self, text: str) -> Iterator[tuple[str, np.ndarray]]:
        """按句合成，逐句 yield (句子, 音频)。

        调用方拿到第一句就能开播，不必等整段——
        这正是 VITS 非自回归下唯一的低延迟路径。
        """
        for seg in split_sentences(text):
            yield seg, self.synthesize(seg)

    def synthesize_all(self, text: str) -> np.ndarray:
        """合成整段并拼接。用于离线/测试，不适合实时播放。"""
        chunks = [a for _, a in self.synthesize_sentences(text) if len(a)]
        if not chunks:
            return np.zeros(0, dtype=np.float32)
        return np.concatenate(chunks)


__all__ = [
    "LocalTTS", "TTSError", "split_sentences", "text_for_speech",
    "MAX_CHARS_PER_SEGMENT",
]