import requests
import argparse
import os
import time

# 默认服务地址
BASE_URL = "http://127.0.0.1:8000"

def test_health():
    """检查服务是否在线"""
    resp = requests.get(f"{BASE_URL}/health")
    print(f"Health check: {resp.status_code} {resp.json()}")

def test_tts(text, speaker=None, speed=1.0, output="output.wav"):
    """
    调用 TTS 接口并保存音频文件
    """
    params = {"text": text, "speed": speed}
    if speaker:
        params["speaker"] = speaker

    print(f"Requesting TTS...")
    start = time.time()
    resp = requests.get(f"{BASE_URL}/tts", params=params)
    elapsed = time.time() - start

    if resp.status_code == 200:
        with open(output, "wb") as f:
            f.write(resp.content)
        print(f"Saved audio to: {output}")
        print(f"Audio size: {os.path.getsize(output)} bytes")
        print(f"Time taken: {elapsed:.2f}s")
    else:
        print(f"Error {resp.status_code}: {resp.text}")

def test_speakers():
    """尝试获取可用说话人列表（通过 health 端点里可能不返回，这里直接尝试几个常见说话人）"""
    # MeloTTS 中文常见说话人: ZH, ZH_CN, ZH_TW 等，我们可以列出文档或手动尝试
    test_speakers = ["ZH", "ZH_CN", "ZH_TW", "EN", "JP"]
    print("Testing available speakers...")
    for spk in test_speakers:
        params = {"text": "测试", "speaker": spk}
        resp = requests.get(f"{BASE_URL}/tts", params=params)
        if resp.status_code == 200:
            print(f"  {spk}: OK")
        else:
            print(f"  {spk}: {resp.text[:50]}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test MeloTTS API server")
    parser.add_argument("--url", default=BASE_URL, help="Server base URL")
    parser.add_argument("--text", default="你好，这是一个语音合成测试。", help="Text to synthesize")
    parser.add_argument("--speaker", default=None, help="Speaker ID (e.g., ZH)")
    parser.add_argument("--speed", type=float, default=1.0, help="Speaking speed")
    parser.add_argument("--output", default="test_output.wav", help="Output WAV filename")
    parser.add_argument("--speakers", action="store_true", help="List available speakers")

    args = parser.parse_args()
    BASE_URL = args.url

    # 先检查服务可用性
    try:
        test_health()
    except requests.ConnectionError:
        print(f"Cannot connect to server at {BASE_URL}. Make sure the FastAPI app is running.")
        exit(1)

    if args.speakers:
        test_speakers()
    else:
        test_tts(args.text, args.speaker, args.speed, args.output)