#!/usr/bin/env python3
"""
Genie TTS 极速服务 - 支持文本分块并行合成 + 线程池预热
适配 ARM 大小核架构（Cortex-X925 + Cortex-A725, 20核），推荐 MAX_WORKERS=8
"""

import os
os.environ["GENIE_DATA_DIR"] = r"/home/admin/minor_agent/modelscope/hub/models/genie_tts/GenieData"

import io
import tempfile
import re
import wave
import multiprocessing
from concurrent.futures import ThreadPoolExecutor, as_completed
from contextlib import asynccontextmanager

from fastapi import FastAPI, Query
from fastapi.responses import Response
import genie_tts as genie

# ======================== 配置区域 ========================
DEFAULT_CHARACTER = "minor"
DEFAULT_ONNX_MODEL_DIR = r"/home/admin/minor_agent/modelscope/models/genie_tts/CharacterModels/v2ProPlus/feibi/tts_models"
DEFAULT_REFERENCE_AUDIO = "/home/admin/minor_agent/modelscope/models/genie_tts/CharacterModels/v2ProPlus/minor/minor.wav"
DEFAULT_REFERENCE_TEXT = "你好你好你好，这是一段...测试录音~"
DEFAULT_LANGUAGE = "zh"
TMP_DIR = "/dev/shm"               # 内存文件系统，减少磁盘 I/O

# 分块与并行参数
MIN_EFFECTIVE_CHARS = 5            # 每个块至少包含的有效字符数
MAX_WORKERS = 4                    # 建议 4~8，ARM 大小核不要超过大核数（10）
POOL_WARMUP = True                 # 是否在线程池创建后立即预热
# =========================================================

# ---------- 全局线程池（应用生命周期内复用） ----------
tts_pool: ThreadPoolExecutor = None

# ---------- 正则工具 ----------
SENT_SPLIT_PATTERN = re.compile(r'(?<=[。！？；，、])')
REMOVE_PUNCT_PATTERN = re.compile(r'[。！？；，、：；“”（）《》【】\-\_\—\~\.\,\!\?\s·]')

def effective_length(text: str) -> int:
    """计算有效字符数（去除标点、连接符、空格）"""
    return len(REMOVE_PUNCT_PATTERN.sub('', text))

def split_text(text: str) -> list[str]:
    """将文本切分为多个语义块"""
    raw_chunks = [c.strip() for c in SENT_SPLIT_PATTERN.split(text) if c.strip()]
    if not raw_chunks:
        return [text]

    merged = []
    for chunk in raw_chunks:
        if merged and effective_length(chunk) < MIN_EFFECTIVE_CHARS:
            merged[-1] += chunk
        else:
            merged.append(chunk)

    if len(merged) == 1:
        return merged

    if effective_length(merged[-1]) < MIN_EFFECTIVE_CHARS and len(merged) > 1:
        merged[-2] += merged[-1]
        merged.pop()

    return merged

# ---------- 音频拼接 ----------
def merge_wavs(file_paths: list[str]) -> bytes:
    """按路径顺序拼接多个 WAV 文件"""
    if len(file_paths) == 1:
        with open(file_paths[0], 'rb') as f:
            return f.read()

    data_parts = []
    params = None
    for path in file_paths:
        with wave.open(path, 'rb') as wf:
            if params is None:
                params = wf.getparams()
            data_parts.append(wf.readframes(wf.getnframes()))

    buf = io.BytesIO()
    with wave.open(buf, 'wb') as wf_out:
        wf_out.setparams(params)
        for frames in data_parts:
            wf_out.writeframes(frames)
    return buf.getvalue()

# ---------- 预热函数 ----------
def _warmup_task(idx: int):
    """空任务用于预热线程"""
    pass

def warmup_pool(pool: ThreadPoolExecutor, num_workers: int):
    """向线程池提交空任务以预热所有线程"""
    futures = []
    for i in range(num_workers):
        futures.append(pool.submit(_warmup_task, i))
    # 等待所有预热任务完成
    for f in futures:
        f.result()
    print(f"✅ 线程池预热完成 (workers={num_workers})")

# ---------- 并行 TTS（使用全局线程池） ----------
def parallel_tts(chunks: list[str]) -> list[str]:
    """
    多线程并行生成所有语音块，返回按原始顺序的临时文件路径列表。
    """
    tmp_paths = []
    for _ in chunks:
        tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False, dir=TMP_DIR)
        tmp.close()
        tmp_paths.append(tmp.name)

    try:
        # 使用全局预热的线程池
        future_to_idx = {}
        for idx, chunk in enumerate(chunks):
            future = tts_pool.submit(
                genie.tts,
                character_name=DEFAULT_CHARACTER,
                text=chunk,
                play=False,
                save_path=tmp_paths[idx]
            )
            future_to_idx[future] = idx

        for future in as_completed(future_to_idx):
            future.result()   # 传递异常
    except:
        for p in tmp_paths:
            if os.path.exists(p):
                os.unlink(p)
        raise
    return tmp_paths

# ---------- 应用生命周期 ----------
@asynccontextmanager
async def lifespan(app: FastAPI):
    global tts_pool
    
    print("🚀 正在启动 Genie TTS 极简服务...")
    try:
        # 1. 加载模型
        genie.load_character(
            character_name=DEFAULT_CHARACTER,
            onnx_model_dir=DEFAULT_ONNX_MODEL_DIR,
            language=DEFAULT_LANGUAGE
        )
        genie.set_reference_audio(
            character_name=DEFAULT_CHARACTER,
            audio_path=DEFAULT_REFERENCE_AUDIO,
            audio_text=DEFAULT_REFERENCE_TEXT,
            language=DEFAULT_LANGUAGE
        )
        
        # 2. 创建并预热全局线程池
        print(f"🔥 正在创建线程池并预热 (max_workers={MAX_WORKERS})...")
        tts_pool = ThreadPoolExecutor(max_workers=MAX_WORKERS)
        if POOL_WARMUP:
            warmup_pool(tts_pool, MAX_WORKERS)
        
        print("✅ Genie TTS 服务启动成功！")
        print(f"   - 线程池: max_workers={MAX_WORKERS}, 已预热")
        print("👉 调用方式：http://localhost:8000/tts?text=你要说的话")
    except Exception as e:
        print(f"❌ 服务启动失败: {e}")
        raise e
    
    yield
    
    print("\n🛑 正在关闭服务...")
    # 清理线程池
    if tts_pool:
        print("   关闭线程池...")
        tts_pool.shutdown(wait=True)
        tts_pool = None
    
    # 卸载模型
    try:
        genie.unload_character(DEFAULT_CHARACTER)
        print("✅ 资源已释放")
    except:
        pass

app = FastAPI(lifespan=lifespan)

# ---------- 极简 API ----------
@app.get("/tts")
async def text_to_speech(text: str = Query(..., description="待合成的文本")):
    if not text.strip():
        return Response(content="文本不能为空", status_code=400, media_type="text/plain")

    chunks = split_text(text)

    try:
        if len(chunks) == 1:
            # 单块直接合成，避免多线程开销
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False, dir=TMP_DIR) as tmp:
                tmp_path = tmp.name
            try:
                genie.tts(
                    character_name=DEFAULT_CHARACTER,
                    text=chunks[0],
                    play=False,
                    save_path=tmp_path
                )
                with open(tmp_path, "rb") as f:
                    audio_bytes = f.read()
                return Response(content=audio_bytes, media_type="audio/wav")
            finally:
                if os.path.exists(tmp_path):
                    os.unlink(tmp_path)
        else:
            # 多块并行合成（使用预热好的线程池）
            tmp_paths = parallel_tts(chunks)
            try:
                merged_audio = merge_wavs(tmp_paths)
                return Response(content=merged_audio, media_type="audio/wav")
            finally:
                for p in tmp_paths:
                    if os.path.exists(p):
                        os.unlink(p)
    except Exception as e:
        return Response(content=f"生成失败: {str(e)}", status_code=500, media_type="text/plain")

@app.get("/health")
async def health():
    return {
        "status": "ok",
        "pool_active": tts_pool is not None,
        "max_workers": MAX_WORKERS
    }

# ---------- 启动入口 ----------
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)