#!/usr/bin/env python3
"""Qwen3-TTS 服务（独立进程）。

为什么要单独起服务，而不是在助手进程里 import：
    qwen-tts 和 vllm-omni 的依赖互斥——前者要 transformers 4.57 + hf_hub 0.36，
    后者要 5.14 + 1.31。装在一个 venv 里会互相破坏（实测把 vllm-omni 搞崩过）。
    所以分成两个进程，用 HTTP 通信。

接口按 **OpenAI 的 /v1/audio/speech 形状**实现：
    - 助手那边换 TTS 后端（本地服务 / 云端 API）时调用方不用改
    - 多出来的 `instructions` 字段是 Qwen3-TTS 的风格指令，
      OpenAI 也有同名字段，语义一致

实测数据（RTX 3090，1.7B，非流式）：
    RTF 0.55–0.73，输出 24kHz 单声道
    flash-attn 装不上（预编译包最高 torch 2.9，本机是 2.14），
    所以走的是 PyTorch 手写注意力，比装了 flash-attn 慢。

运行：
    ~/tts-qwen/.venv/bin/python tts_server.py --port 18040
"""

from __future__ import annotations

import argparse
import asyncio
import io
import logging
import os
import time
import threading
from contextlib import asynccontextmanager
from pathlib import Path

import numpy as np
import soundfile as sf
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel, Field

log = logging.getLogger("tts")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

# ---------------------------------------------------------------- 默认音色

#: 用户试听选定的组合（见 out/tts_qwen/）：Serena 音色 + 轻快活泼的语气。
#: 时长比对确认：instruct_7849.wav 是 1.44s，对应「用轻快活泼的语气说」。
DEFAULT_VOICE = os.environ.get("TTS_DEFAULT_VOICE", "Serena")
DEFAULT_INSTRUCT = os.environ.get("TTS_DEFAULT_INSTRUCT", "用轻快活泼的语气说")
DEFAULT_LANGUAGE = os.environ.get("TTS_DEFAULT_LANGUAGE", "Chinese")

#: 服务端支持的音色（来自模型 README）
VOICES = {
    "Vivian": "明亮微冷的年轻女声",
    "Serena": "温暖柔和的年轻女声",
    "Uncle_Fu": "低沉醇厚的成熟男声",
    "Dylan": "清朗自然的京腔男声",
    "Eric": "略带沙哑的成都男声",
    "Ryan": "富有节奏感的男声（英文）",
    "Aiden": "阳光美式男声（英文）",
    "Ono_Anna": "轻快俏皮的日语女声",
    "Sohee": "温暖富有情感的韩语女声",
}

_model = None
#: 模型是单实例、占用 GPU，并发调用必须串行化，否则会互相踩
_lock = threading.Lock()


class SpeechRequest(BaseModel):
    """OpenAI /v1/audio/speech 的形状，加一个 Qwen 特有的语言字段。"""
    input: str = Field(..., description="要合成的文本")
    voice: str = Field(DEFAULT_VOICE, description="音色名，见 /v1/voices")
    instructions: str | None = Field(DEFAULT_INSTRUCT, description="自然语言风格指令")
    language: str = Field(DEFAULT_LANGUAGE, description="语言，中文用 Chinese")
    speed: float = Field(1.0, description="语速（本模型通过 instructions 控制，此项仅保留兼容）")
    response_format: str = Field("wav", description="输出格式，目前只支持 wav")


@asynccontextmanager
async def lifespan(app: FastAPI):
    global _model
    model_dir = os.environ.get("QWEN_TTS_MODEL",
                               os.path.expanduser("~/tts-service/models/Qwen3-TTS-1.7B-CustomVoice"))
    log.info("加载模型 %s ...", model_dir)
    t0 = time.perf_counter()
    from qwen_tts import Qwen3TTSModel
    # 在线程里加载，别阻塞事件循环
    _model = await asyncio.to_thread(
        Qwen3TTSModel.from_pretrained, model_dir, device_map="cuda:0")
    log.info("模型就绪，耗时 %.1fs", time.perf_counter() - t0)
    yield
    _model = None


app = FastAPI(title="Qwen3-TTS Service", lifespan=lifespan)


@app.get("/health")
async def health():
    return {"status": "ok" if _model is not None else "loading",
            "model_loaded": _model is not None,
            "default_voice": DEFAULT_VOICE}


@app.get("/v1/voices")
async def voices():
    return {"default": DEFAULT_VOICE, "default_instructions": DEFAULT_INSTRUCT,
            "voices": VOICES}


def _synth(text: str, voice: str, instruct: str | None, language: str) -> tuple[np.ndarray, int]:
    """同步合成。调用方负责串行化。"""
    kwargs = {"text": text, "language": language, "speaker": voice}
    if instruct:
        kwargs["instruct"] = instruct
    wavs, sr = _model.generate_custom_voice(**kwargs)
    # 模型可能按内部切分返回多段，拼起来
    if isinstance(wavs, list):
        audio = np.concatenate([np.asarray(w, dtype=np.float32) for w in wavs])
    else:
        audio = np.asarray(wavs, dtype=np.float32)
    return audio, sr


@app.post("/v1/audio/speech")
async def speech(req: SpeechRequest):
    if _model is None:
        raise HTTPException(503, "模型尚未加载完成")
    text = (req.input or "").strip()
    if not text:
        raise HTTPException(400, "input 不能为空")
    if req.voice not in VOICES:
        raise HTTPException(400, f"未知音色 {req.voice!r}，可用：{sorted(VOICES)}")

    t0 = time.perf_counter()
    try:
        # 模型单实例占 GPU，串行执行；放线程里避免阻塞事件循环
        audio, sr = await asyncio.to_thread(
            lambda: _with_lock(_synth, text, req.voice, req.instructions, req.language))
    except Exception as e:  # noqa: BLE001
        log.exception("合成失败")
        raise HTTPException(500, f"合成失败: {type(e).__name__}: {e}") from e

    gen_ms = (time.perf_counter() - t0) * 1000
    dur_s = len(audio) / sr
    rtf = (gen_ms / 1000) / dur_s if dur_s else 0
    log.info("合成 %d 字 → %.2fs 音频，耗时 %.0fms（RTF %.2f）",
             len(text), dur_s, gen_ms, rtf)

    buf = io.BytesIO()
    sf.write(buf, audio, sr, format="WAV")
    return Response(
        content=buf.getvalue(),
        media_type="audio/wav",
        headers={
            "X-Sample-Rate": str(sr),
            "X-Duration-Seconds": f"{dur_s:.3f}",
            "X-Gen-Ms": f"{gen_ms:.0f}",
            "X-RTF": f"{rtf:.3f}",
        },
    )


def _with_lock(fn, *args):
    with _lock:
        return fn(*args)


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=18040)
    a = p.parse_args()
    uvicorn.run(app, host=a.host, port=a.port, log_level="info")