#!/usr/bin/env python3
"""
VoxCPM 流式 TTS 服务 —— 按句分块合成，SSE 流式推送 PCM 音频。

使用 OpenBMB VoxCPM 模型（Whispera 同款），无需 Genie TTS。
配合前端 Web Audio API 实现边合成边播放，无需保存音频文件。

启动方式:
    python streaming_tts_server.py --port 8902 --model-path ./models/VoxCPM
"""

import os
import sys
import json
import base64
import time
import threading
import argparse
import warnings
from contextlib import asynccontextmanager
from pathlib import Path

# 抑制 PyTorch weight_norm 弃用警告（来自 voxcpm 内部）
warnings.filterwarnings("ignore", message=".*weight_norm.*deprecated.*")

# ======================== 配置区域 ========================
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# 启用 TF32 加速（Ampere+ GPU）
import torch
torch.set_float32_matmul_precision("high")

# 确保 torchaudio 已安装（VoxCPM 参考音频需要）
try:
    import torchaudio  # noqa: F401
except ImportError:
    pass  # 无参考音频时可选

# 默认参数（可通过命令行覆盖）
DEFAULT_MODEL_PATH = os.path.join(BASE_DIR, "models", "VoxCPM1__5")
DEFAULT_PORT = 8902

# 参考音频（用于统一音色，同一人物说的一段话 + 对应文本）
# 如果配置了这对参数，所有合成都会克隆该音色，各句音色保持一致
DEFAULT_PROMPT_WAV = os.path.join(BASE_DIR, "reference_audio.wav")
DEFAULT_PROMPT_TEXT = "日子缓缓前行，不必急于追赶，照顾好自己的心情，认真过好当下的每一刻就足够了。"

# 音频质量参数
DEFAULT_CFG_VALUE = 1.5          # 音色贴合度 (1.0~3.0)，越高越贴近参考音频但可能不自然
DEFAULT_INFERENCE_STEPS = 15     # 推理步数 (4~30)，越高音质越好但越慢
# =========================================================

# 请求级取消标志
_cancel_flags: dict[str, threading.Event] = {}
_cancel_lock = threading.Lock()

# 当前活跃请求 ID（新请求到达时自动取消旧请求）
_current_request_id: str | None = None
_current_req_lock = threading.Lock()

# 模型锁（VoxCPM 只允许单线程使用，新请求取消旧请求后等待锁释放）
_model_lock = threading.Lock()


def set_cancel(request_id: str) -> None:
    with _cancel_lock:
        evt = _cancel_flags.get(request_id)
        if evt is None:
            evt = threading.Event()
            _cancel_flags[request_id] = evt
        evt.set()


def is_cancelled(request_id: str) -> bool:
    with _cancel_lock:
        evt = _cancel_flags.get(request_id)
        return evt is not None and evt.is_set()


def cancel_current_and_set(new_request_id: str) -> str | None:
    """取消当前活跃请求并注册新请求 ID，返回被取消的旧请求 ID。"""
    global _current_request_id
    with _current_req_lock:
        old = _current_request_id
        _current_request_id = new_request_id
    if old and old != new_request_id:
        set_cancel(old)
        # 若旧请求持有模型锁，强制释放以让新请求立即获取
        try:
            if _model_lock.locked():
                _model_lock.release()
        except RuntimeError:
            pass
        print(f"[VoxCPM TTS] 打断旧请求 {old}，开始新请求 {new_request_id}")
    return old


# Lazy-loaded model singleton
_model = None
_singleton_lock = threading.Lock()
_model_sample_rate = 24000


def _get_model(model_path: str):
    """延迟加载 VoxCPM 模型（线程安全）。"""
    global _model, _model_sample_rate
    if _model is not None:
        return _model, _model_sample_rate

    with _singleton_lock:
        if _model is not None:
            return _model, _model_sample_rate

        if "voxcpm" not in sys.modules:
            try:
                from voxcpm import VoxCPM  # type: ignore
            except ImportError:
                raise ImportError(
                    "无法导入 voxcpm 模块。请确保 Whispera-main 项目的 "
                    "voxcpm-tts-streaming-module/src 在 sys.path 中，"
                    "或 pip install 了 voxcpm 包。"
                )

        from voxcpm import VoxCPM  # type: ignore
        mp = _resolve_model_path(model_path)
        print(f"[VoxCPM TTS] 加载模型: {mp}")
        _model = VoxCPM(
            voxcpm_model_path=mp,
            enable_denoiser=False,
            optimize=True,
        )
        _model_sample_rate = int(_model.tts_model.sample_rate)
        print(f"[VoxCPM TTS] 模型加载成功，采样率: {_model_sample_rate} Hz")
        return _model, _model_sample_rate


def _resolve_model_path(model_path: str) -> str:
    """解析模型路径：支持多种命名格式。"""
    base = Path(model_path).expanduser().resolve()
    if not base.exists():
        raise FileNotFoundError(f"VoxCPM 模型路径不存在: {base}")

    # 1. 直接就是模型目录（有 config.json）
    if (base / "config.json").is_file():
        return str(base)

    # 2. 检查常见子目录名
    for sub in ["openbmb__VoxCPM2", "openbmb__VoxCPM1.5",
                "VoxCPM2", "VoxCPM1.5", "VoxCPM1__5"]:
        candidate = base / sub
        if candidate.is_dir() and (candidate / "config.json").is_file():
            return str(candidate)

    raise FileNotFoundError(
        f"在 {base} 下未找到 VoxCPM 模型（需要包含 config.json 的子目录，"
        f"如 openbmb__VoxCPM1.5/、VoxCPM1__5/ 等）"
    )


def _build_sse(data: dict) -> str:
    return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"


def _stream_audio_events(text: str, request_id: str, model_path: str):
    """生成 SSE 音频流事件（同步，在独立线程中运行）。"""
    import numpy as np

    try:
        model, sample_rate = _get_model(model_path)
    except Exception as e:
        yield _build_sse({"type": "audio.error", "request_id": request_id, "message": str(e)})
        return

    # 开始
    yield _build_sse({
        "type": "audio.start",
        "request_id": request_id,
        "sample_rate": sample_rate,
        "audio_format": "pcm_f32le",
    })

    chunk_index = 0
    total_samples = 0

    # 获取模型锁（前一请求已被 cancel_current_and_set 取消，锁应已释放）
    # 短超时 + 重试，避免用户等待
    for attempt in range(3):
        if _model_lock.acquire(timeout=3):
            break
        if attempt < 2:
            print(f"[VoxCPM TTS] 等待模型锁释放... (尝试 {attempt + 1}/3)")
    else:
        yield _build_sse({"type": "audio.error", "request_id": request_id,
                          "message": "等待模型超时，请重试"})
        return
    try:
        gen = model.generate_streaming(
            text=text,
            prompt_wav_path=DEFAULT_PROMPT_WAV or None,
            prompt_text=DEFAULT_PROMPT_TEXT or None,
            cfg_value=DEFAULT_CFG_VALUE,
            inference_timesteps=DEFAULT_INFERENCE_STEPS,
            retry_badcase=False,
        )

        for audio_chunk in gen:
            if is_cancelled(request_id):
                yield _build_sse({"type": "audio.cancelled", "request_id": request_id})
                return

            arr = np.asarray(audio_chunk, dtype=np.float32).ravel()
            pcm_bytes = np.ascontiguousarray(arr, dtype=np.float32).tobytes()
            num_samples = arr.size
            total_samples += num_samples
            b64_data = base64.b64encode(pcm_bytes).decode("ascii")

            yield _build_sse({
                "type": "audio.chunk",
                "request_id": request_id,
                "index": chunk_index,
                "num_samples": num_samples,
                "sample_rate": sample_rate,
                "data": b64_data,
            })
            chunk_index += 1
    finally:
        _model_lock.release()

    yield _build_sse({
        "type": "audio.done",
        "request_id": request_id,
        "total_samples": total_samples,
        "interrupted": False,
    })


# ======================== FastAPI ========================

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse


@asynccontextmanager
async def lifespan(app: FastAPI):
    mp = _model_args.get("model_path", DEFAULT_MODEL_PATH)
    print(f"[VoxCPM TTS] 服务启动中，模型路径: {mp}")
    print(f"[VoxCPM TTS] 调用方式: POST /stream_tts body: {{'text': '...', 'request_id': '...'}}")
    # 启动时预加载模型（避免首次请求等待）
    try:
        _get_model(mp)
    except Exception as e:
        print(f"[VoxCPM TTS] 模型预加载失败: {e}")
    yield
    print("[VoxCPM TTS] 服务关闭")


app = FastAPI(title="VoxCPM Streaming TTS", lifespan=lifespan)


@app.post("/stream_tts")
async def stream_tts(request: Request):
    """流式 TTS 端点：接收文本，SSE 流式返回 PCM f32le 音频块。"""
    try:
        body = await request.json()
    except Exception:
        form = await request.form()
        body = dict(form)

    text = (body.get("text") or "").strip()
    request_id = body.get("request_id") or f"tts-{int(time.time() * 1000)}"

    if not text:
        return StreamingResponse(
            iter([_build_sse({"type": "audio.error", "request_id": request_id, "message": "文本为空"})]),
            media_type="text/event-stream",
        )

    # 新请求到达：自动取消当前活跃的旧请求
    cancel_current_and_set(request_id)

    return StreamingResponse(
        _stream_audio_events(text, request_id, _model_args.get("model_path", DEFAULT_MODEL_PATH)),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@app.post("/cancel_tts")
async def cancel_tts(request: Request):
    """取消指定 request_id 的 TTS 合成。"""
    try:
        body = await request.json()
    except Exception:
        form = await request.form()
        body = dict(form)

    request_id = body.get("request_id", "")
    if request_id:
        set_cancel(request_id)
        return {"status": "cancelled", "request_id": request_id}
    return {"status": "error", "message": "缺少 request_id"}


@app.get("/health")
async def health():
    return {"status": "ok"}


# ======================== 启动入口 ========================
_model_args: dict = {}

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="VoxCPM Streaming TTS Server")
    parser.add_argument("--port", type=int, default=DEFAULT_PORT, help=f"服务端口（默认 {DEFAULT_PORT}）")
    parser.add_argument("--model-path", type=str, default=DEFAULT_MODEL_PATH, help="VoxCPM 模型目录")
    args = parser.parse_args()

    _model_args["model_path"] = args.model_path

    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=args.port)