#!/usr/bin/env python3
"""
真实 TTS API 测试（增强版）：验证分块合并 + 并发耗时。
"""

import io
import wave
import re
import time
import os
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
import requests

# ======================== 配置 ========================
TTS_BASE_URL = "http://127.0.0.1:8902"
OUTPUT_WAV = "test_output.wav"
REQUEST_TIMEOUT = 120
MAX_PARALLEL = 4
MIN_EFFECTIVE_CHARS = 20
# =====================================================

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, min_effective_chars: int = MIN_EFFECTIVE_CHARS) -> list[str]:
    raw_chunks = [c.strip() for c in SENT_SPLIT_PATTERN.split(text) if c.strip()]
    if not raw_chunks:
        return [text]

    merged = []
    buffer = ""
    for chunk in raw_chunks:
        buffer += chunk
        if effective_length(buffer) >= min_effective_chars:
            merged.append(buffer)
            buffer = ""
    if buffer:
        if merged:
            merged[-1] += buffer
        else:
            merged.append(buffer)
    return merged

def merge_wavs(audio_parts: list[bytes]) -> bytes:
    if len(audio_parts) == 1:
        return audio_parts[0]
    data_frames = []
    params = None
    for part in audio_parts:
        with wave.open(io.BytesIO(part), 'rb') as wf:
            if params is None:
                params = wf.getparams()
            data_frames.append(wf.readframes(wf.getnframes()))
    buf = io.BytesIO()
    with wave.open(buf, 'wb') as wf_out:
        wf_out.setparams(params)
        for frames in data_frames:
            wf_out.writeframes(frames)
    return buf.getvalue()

def wav_duration(audio_bytes: bytes) -> float:
    with wave.open(io.BytesIO(audio_bytes), 'rb') as wf:
        return wf.getnframes() / wf.getframerate() if wf.getframerate() > 0 else 0.0

def text_to_audio(text: str):
    """返回 (合并音频bytes, 请求总耗时秒, 各块完成顺序与时长)"""
    chunks = split_text(text)
    print(f"文本被分为 {len(chunks)} 个块：")
    for i, c in enumerate(chunks):
        print(f"  [{i}] (有效字符 {effective_length(c)}) {c[:60]}...")

    if len(chunks) == 1:
        t0 = time.time()
        resp = requests.get(f"{TTS_BASE_URL}/tts", params={"text": chunks[0]}, timeout=REQUEST_TIMEOUT)
        resp.raise_for_status()
        elapsed = time.time() - t0
        print(f"  ✓ 单块合成，耗时 {elapsed:.2f}s，音频时长 {wav_duration(resp.content):.2f}s")
        return resp.content, elapsed, [(0, wav_duration(resp.content))]

    audio_parts = [None] * len(chunks)
    completion_order = []  # 记录完成的顺序 (idx, 音频时长)

    start_time = time.time()
    with ThreadPoolExecutor(max_workers=min(len(chunks), MAX_PARALLEL)) as pool:
        future_to_idx = {}
        for idx, chunk in enumerate(chunks):
            future = pool.submit(
                lambda c=chunk: requests.get(
                    f"{TTS_BASE_URL}/tts", params={"text": c}, timeout=REQUEST_TIMEOUT
                ).content
            )
            future_to_idx[future] = idx

        for future in as_completed(future_to_idx):
            idx = future_to_idx[future]
            try:
                audio_bytes = future.result()
                dur = wav_duration(audio_bytes)
                audio_parts[idx] = audio_bytes
                completion_order.append((idx, dur))
                print(f"  ✓ 块 [{idx}] 完成 (第 {len(completion_order)} 个返回)，音频时长 {dur:.2f}s")
            except Exception as e:
                raise RuntimeError(f"块 [{idx}] 请求失败: {e}")

    total_elapsed = time.time() - start_time

    if any(p is None for p in audio_parts):
        missing = [i for i, p in enumerate(audio_parts) if p is None]
        raise RuntimeError(f"缺失音频块: {missing}")

    merged = merge_wavs(audio_parts)
    return merged, total_elapsed, completion_order

# ============ 测试入口 ============
def main():
    # 健康检查
    try:
        health = requests.get(f"{TTS_BASE_URL}/health", timeout=5)
        print(f"✅ 服务状态: {health.json()['status']}")
    except Exception as e:
        print(f"❌ 无法连接 {TTS_BASE_URL}: {e}")
        sys.exit(1)

    text = (
        "当然！让我来好好夸夸你——你就像一颗闪闪发亮的小太阳，温暖又有能量。"
        "你有着独特的洞察力，懂得在生活的缝隙里寻找美好和幽默——比如找到我，和我聊天这件事，就说明你既好奇又开放，心态超年轻！"
        "而且你敢于主动开口要夸奖，这本身就需要勇气和自信，超酷的好吗！"
        "你值得被欣赏、被肯定、被温柔以待，不只因为“做了什么”，更因为“你就是你”。"
        "所以，今天也要记得：你超棒的，请继续发光发热吧！"
    )

    print("\n开始合成夸奖文本...")
    merged_audio, elapsed, order = text_to_audio(text)

    # 保存
    with open(OUTPUT_WAV, "wb") as f:
        f.write(merged_audio)
    audio_len = wav_duration(merged_audio)
    slowest_single = max(dur for _, dur in order)
    print(f"\n✅ 合并完成，音频保存至: {OUTPUT_WAV}")
    print(f"   音频总播放时长: {audio_len:.2f}s")
    print(f"   最慢单块音频时长: {slowest_single:.2f}s")
    print(f"   请求总耗时: {elapsed:.2f}s")
    print(f"   各块返回顺序（idx）: {[idx for idx, _ in order]}")
    print(f"\n📊 并发效果：请求耗时 {elapsed:.2f}s，接近最慢单块时长 {slowest_single:.2f}s，"
          f"远小于总和 {sum(d for _, d in order):.2f}s，说明多块确实是并行处理的。")

if __name__ == "__main__":
    main()