"""TTS（语音合成）sidecar — Kokoro-82M（ONNX），零 torch。

与 `packages/stt/stt_server.py` 同构：Node 侧分配端口后用 `--port` 传入、靠 `/health`
判断就绪、`--parent-pid` 看门狗防孤儿进程；父进程退出后本进程自杀。
权重**不由本进程下载**（Node 侧 `ttsDownload.ts` 负责，含国内镜像选择），本进程只从
models_dir 加载。

环境变量：
  DSC_TTS_PORT        监听端口（默认 0 = 系统分配）
  DSC_TTS_MODELS_DIR  权重目录（默认 ~/.agent/tts-models）
  DSC_TTS_MODEL       模型文件名（默认 model_uint8f16.onnx）
  DSC_TTS_VOICES      音色文件名（默认 voices-v1.0.bin）

端点：
  GET  /health      进程就绪（不依赖模型加载完成）
  GET  /status      依赖与权重就绪情况（供「缺什么 + 一键装」文案使用）
  POST /synthesize  {"text","voice","speed","lang","stream"}
                    非流式 -> application/wav 整段返回
                    流式   -> audio/L16 裸 PCM（s16le 24kHz 单声道），HTTP/1.0 连接关闭定界

中文音素**必须**走 misaki（见 docs/architecture-design/PERSONAL_REMOTE_AGENT_PLAN.md
gotcha 11）：kokoro-onnx 自带的 `phonemize(text, "cmn")` 走 espeak，产出丢声母且**完全无音调**，
中文因此听不懂——而且不报错、不掉字，只有人耳能发现。Kokoro 的中文音色是用 misaki 的
拼音 IPA（带 `↓→↗↘` 调号）训练的，所以这里自己产出音素，再以 `is_phonemes=True` 送入。
"""

from __future__ import annotations

import asyncio
import io
import json
import os
import re
import sys
import threading
import time
import wave
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any, Callable

import numpy as np

DEFAULT_MODELS_DIR = str(Path.home() / ".agent" / "tts-models")
SAMPLE_RATE = 24_000
DEFAULT_MODEL = "model_uint8f16.onnx"
DEFAULT_VOICES = "voices-v1.0.bin"
# 每档语言的默认音色（Kokoro v1.0 音色名；未知音色会被 resolve_voice 挡下并列出可选值）
DEFAULT_VOICE = {"zh": "zf_xiaoxiao", "en": "af_heart"}
FALLBACK_VOICE = "af_heart"
MAX_TEXT_CHARS = 5_000
# 中英混排切分：英文段（含文件名/路径里的 . _ - /）与其余段。见 gotcha 11 的切分理由。
_MIXED_RE = re.compile(r"[A-Za-z][A-Za-z0-9_.\-/]*|[^A-Za-z]+")


def log(message: str) -> None:
    print(f"[tts] {message}", flush=True)


def _arg_value(name: str) -> str | None:
    """从 sys.argv 取 `--x value` 或 `--x=value`（与 stt_server.py 同一约定）。"""
    argv = sys.argv
    for index, item in enumerate(argv):
        if item == name and index + 1 < len(argv):
            return argv[index + 1]
        if item.startswith(f"{name}="):
            return item.split("=", 1)[1]
    return None


def is_chinese(lang: str) -> bool:
    normalized = (lang or "").lower().replace("_", "-")
    return normalized.startswith("zh") or normalized.startswith("cmn") or normalized.startswith("yue")


def lang_key(lang: str) -> str:
    return "zh" if is_chinese(lang) else "en"


def parent_alive(pid: int) -> bool:
    """父进程是否还在。

    **不要**用 `os.kill(pid, 0)` 探活：Windows 上它并非「发 0 号信号」，CPython 会走
    `OpenProcess(PROCESS_ALL_ACCESS)` + `TerminateProcess`（文档：除 CTRL_C_EVENT /
    CTRL_BREAK_EVENT 外的 sig 会「unconditionally killed by the TerminateProcess API」），
    而且 `PROCESS_ALL_ACCESS` 在子进程带 `CREATE_NO_WINDOW`（Node 的 `windowsHide: true`）
    时**会失败**并抛 OSError → 看门狗误判「父进程已死」→ 服务约 5 秒后自杀，表现为
    第一个请求连接被重置、之后一律 ConnectionRefused。实测见 .agent/tts-probe/shape_probe.py。
    """
    if os.name == "nt":
        import ctypes

        kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
        # 只申请「查询」权限，不做任何写/终止操作
        handle = kernel32.OpenProcess(0x1000, False, pid)
        if not handle:
            return False
        try:
            code = ctypes.c_ulong()
            if not kernel32.GetExitCodeProcess(handle, ctypes.byref(code)):
                return False
            return code.value == 259  # STILL_ACTIVE
        finally:
            kernel32.CloseHandle(handle)
    try:
        os.kill(pid, 0)
        return True
    except (OSError, ValueError):
        return False


def build_phonemes(text: str, lang: str, tokenizer: Any, zh_g2p: Any) -> str:
    """text -> 音素串。非中文交给 kokoro-onnx 原生 espeak 路径；中文走 misaki，
    并把中英混排自己切开——否则英文段会被原样保留、随后被 tokenizer 静默丢弃。"""
    if not is_chinese(lang):
        return tokenizer.phonemize(text, lang)
    parts: list[str] = []
    for match in _MIXED_RE.finditer(text):
        chunk = match.group(0)
        if not chunk.strip():
            continue
        if chunk[0].isascii() and chunk[0].isalpha():
            parts.append(tokenizer.phonemize(chunk, "en-us"))
        else:
            phonemes, _ = zh_g2p(chunk)
            parts.append(phonemes)
    return " ".join(part for part in parts if part)


def to_pcm16(samples: Any) -> bytes:
    clipped = np.clip(np.asarray(samples, dtype=np.float32), -1.0, 1.0)
    return (clipped * 32767.0).astype("<i2").tobytes()


def wav_bytes(samples: Any, sample_rate: int) -> bytes:
    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as handle:
        handle.setnchannels(1)
        handle.setsampwidth(2)
        handle.setframerate(sample_rate)
        handle.writeframes(to_pcm16(samples))
    return buffer.getvalue()


class TtsEngine:
    """懒加载 Kokoro。加载失败时把原因留在 probe() 里，供前端「缺什么」文案使用。"""

    def __init__(self, models_dir: str, model_file: str, voices_file: str) -> None:
        self.models_dir = models_dir
        self.model_file = model_file
        self.voices_file = voices_file
        self.model_path = str(Path(models_dir) / model_file)
        self.voices_path = str(Path(models_dir) / voices_file)
        self._kokoro: Any = None
        self._zh_g2p: Any = None
        self._load_error: str | None = None
        self._lock = threading.Lock()

    def missing_weights(self) -> list[str]:
        pairs = ((self.model_file, self.model_path), (self.voices_file, self.voices_path))
        return [name for name, path in pairs if not Path(path).is_file()]

    def probe(self) -> dict:
        """不加载模型，只探测依赖可导入性与权重是否就位。"""
        try:
            import kokoro_onnx  # noqa: F401
            import misaki  # noqa: F401

            import_error: str | None = None
        except Exception as exc:  # noqa: BLE001
            import_error = f"{type(exc).__name__}: {exc}"
        missing = self.missing_weights()
        status: dict = {
            "installed": import_error is None and not missing,
            "loaded": self._kokoro is not None,
            "model": self.model_file,
            "voicesFile": self.voices_file,
            "voiceNames": self.voice_names(),
            "device": "cpu",
            "missingWeights": missing,
        }
        if import_error:
            status["error"] = import_error
        elif self._load_error:
            status["error"] = self._load_error
        return status

    def load(self) -> Any:
        with self._lock:
            if self._kokoro is not None:
                return self._kokoro
            if self._load_error:
                raise RuntimeError(self._load_error)
            missing = self.missing_weights()
            if missing:
                raise RuntimeError(f"missing weights: {', '.join(missing)}")
            try:
                from kokoro_onnx import Kokoro
                from misaki import zh

                self._zh_g2p = zh.ZHG2P()
                self._kokoro = Kokoro(self.model_path, self.voices_path)
            except Exception as exc:  # noqa: BLE001
                self._load_error = f"{type(exc).__name__}: {exc}"
                raise RuntimeError(self._load_error) from exc
            log(f"model loaded ({self.model_file}, {len(self._kokoro.voices)} voices)")
            return self._kokoro

    def voice_names(self) -> list[str]:
        """音色名清单（voices 文件的 key），**不加载 ONNX 模型**。

        为什么不用 available_voices()：那条路径要过 load()，会把 114 MiB 模型读进内存，
        而 probe() 是被设置页高频轮询的——这个代价付不起。模型已加载时直接取它的键，
        与读文件是同一份数据，只是省一次解析。
        """
        if self._kokoro is not None:
            return sorted(self._kokoro.voices)
        if not Path(self.voices_path).is_file():
            return []
        try:
            import numpy

            return sorted(numpy.load(self.voices_path, allow_pickle=False).files)
        except Exception as exc:  # noqa: BLE001
            # 拿不到清单不该让整个状态接口失败：合成侧仍会按 lang 选默认音色
            log(f"voice list unavailable: {exc}")
            return []

    def available_voices(self) -> list[str]:
        return sorted(self.load().voices)

    def resolve_voice(self, lang: str, voice: str | None) -> str:
        available = self.available_voices()
        wanted = voice or DEFAULT_VOICE.get(lang_key(lang), FALLBACK_VOICE)
        if wanted in available:
            return wanted
        raise ValueError(f"unknown voice {wanted!r}; available: {', '.join(available)}")

    def _phonemes(self, kokoro: Any, text: str, lang: str) -> str:
        return build_phonemes(text, lang, kokoro.tokenizer, self._zh_g2p)

    def synthesize(self, text: str, lang: str, voice: str | None, speed: float) -> tuple[bytes, int]:
        kokoro = self.load()
        phonemes = self._phonemes(kokoro, text, lang)
        if not phonemes.strip():
            raise ValueError("text produced no phonemes")
        audio, rate = kokoro.create(
            phonemes,
            voice=self.resolve_voice(lang, voice),
            speed=speed,
            lang=lang,
            is_phonemes=True,
        )
        return wav_bytes(audio, rate), rate

    def synthesize_stream(
        self, text: str, lang: str, voice: str | None, speed: float, write: Callable[[bytes], None]
    ) -> None:
        """边合成边写裸 PCM。放弃生成器即停止后续批次，客户端断连自然收敛。"""
        kokoro = self.load()
        phonemes = self._phonemes(kokoro, text, lang)
        if not phonemes.strip():
            raise ValueError("text produced no phonemes")
        style = self.resolve_voice(lang, voice)

        async def run() -> None:
            async for audio, _rate in kokoro.create_stream(
                phonemes, voice=style, speed=speed, lang=lang, is_phonemes=True
            ):
                write(to_pcm16(audio))

        asyncio.run(run())


class TtsHttpServer(ThreadingHTTPServer):
    daemon_threads = True

    def __init__(self, address: tuple[str, int], engine: TtsEngine) -> None:
        super().__init__(address, TtsHandler)
        self.engine = engine


class TtsHandler(BaseHTTPRequestHandler):
    server_version = "DeepSeekCodeTts/1.0"
    # 流式响应无 Content-Length，靠连接关闭定界；HTTP/1.0 正好是这个语义
    protocol_version = "HTTP/1.0"

    def log_message(self, fmt: str, *args: Any) -> None:
        """静音 http.server 默认的逐请求 stderr 噪音。"""

    @property
    def engine(self) -> TtsEngine:
        return self.server.engine  # type: ignore[attr-defined]

    def _send_json(self, status: int, payload: dict) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self) -> None:  # noqa: N802 (http.server 的约定命名)
        if self.path == "/health":
            self._send_json(200, {"ok": True})
        elif self.path == "/status":
            self._send_json(200, self.engine.probe())
        else:
            self._send_json(404, {"error": "not found"})

    def do_POST(self) -> None:  # noqa: N802
        if self.path != "/synthesize":
            self._send_json(404, {"error": "not found"})
            return
        try:
            length = int(self.headers.get("Content-Length") or 0)
            payload = json.loads(self.rfile.read(length) or b"{}")
        except Exception as exc:  # noqa: BLE001
            self._send_json(400, {"error": f"bad request body: {exc}"})
            return
        log(f"-> synthesize {len(payload.get('text') or '')} chars lang={payload.get('lang')}")

        text = str(payload.get("text") or "").strip()
        if not text:
            self._send_json(400, {"error": "text is required"})
            return
        if len(text) > MAX_TEXT_CHARS:
            self._send_json(400, {"error": f"text too long (> {MAX_TEXT_CHARS} chars)"})
            return

        lang = str(payload.get("lang") or "en-us")
        voice = payload.get("voice") or None
        speed = float(payload.get("speed") or 1.0)
        if not 0.5 <= speed <= 2.0:
            self._send_json(400, {"error": f"speed must be within 0.5~2.0, got {speed}"})
            return

        if payload.get("stream"):
            self._stream(text, lang, voice, speed)
            return

        try:
            audio, rate = self.engine.synthesize(text, lang, voice, speed)
        except ValueError as exc:
            self._send_json(400, {"error": str(exc)})
        except Exception as exc:  # noqa: BLE001
            self._send_json(500, {"error": f"{type(exc).__name__}: {exc}"})
        else:
            log(f"<- synthesized {len(audio)} bytes, responding")
            self.send_response(200)
            self.send_header("Content-Type", "audio/wav")
            self.send_header("X-Sample-Rate", str(rate))
            self.send_header("Content-Length", str(len(audio)))
            self.end_headers()
            self.wfile.write(audio)

    def _stream(self, text: str, lang: str, voice: str | None, speed: float) -> None:
        """头一旦发出就无法再改状态码，所以失败只能记日志 + 断流（客户端按截断处理）。"""
        self.send_response(200)
        self.send_header("Content-Type", "audio/L16")
        self.send_header("X-Sample-Rate", str(SAMPLE_RATE))
        self.send_header("Connection", "close")
        self.end_headers()
        try:
            self.engine.synthesize_stream(text, lang, voice, speed, self.wfile.write)
        except (BrokenPipeError, ConnectionResetError):
            log("client closed the stream")
        except Exception as exc:  # noqa: BLE001
            log(f"stream failed: {type(exc).__name__}: {exc}")


def main() -> None:
    models_dir = os.environ.get("DSC_TTS_MODELS_DIR") or DEFAULT_MODELS_DIR
    Path(models_dir).mkdir(parents=True, exist_ok=True)
    model_file = os.environ.get("DSC_TTS_MODEL") or DEFAULT_MODEL
    voices_file = os.environ.get("DSC_TTS_VOICES") or DEFAULT_VOICES
    port = int(_arg_value("--port") or os.environ.get("DSC_TTS_PORT") or "0")
    raw_parent = _arg_value("--parent-pid")
    parent_pid = int(raw_parent) if raw_parent and raw_parent.isdigit() else None

    engine = TtsEngine(models_dir, model_file, voices_file)
    server = TtsHttpServer(("127.0.0.1", port), engine)
    log(f"listening on 127.0.0.1:{server.server_address[1]} (models_dir={models_dir})")

    def watchdog() -> None:
        """父进程（Node server）退出后自动退出，避免残留孤儿进程（Windows 上尤为必要）。"""
        while True:
            time.sleep(5)
            if not parent_alive(parent_pid):
                log(f"parent pid {parent_pid} gone, exiting")
                server.shutdown()
                return

    if parent_pid is not None:
        threading.Thread(target=watchdog, daemon=True).start()

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()