#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
TTS/ASR模块集成测试 — MLX/Qwen3-ASR 版本
测试API端点和完整流程（需要运行后端服务）

运行方式:
    pytest backend/tests/test_tts_asr_integration.py -v -s
    python backend/tests/test_tts_asr_integration.py --test asr

MLX 模型通过 ModelScope (aufklarer/Qwen3-ASR) + ForcedAligner
"""

import argparse
import base64
import io
import os
import sys
import time
import unittest
from typing import Optional

try:
    import httpx  # type: ignore
except ImportError:
    print("httpx 未安装，跳过集成测试")
    sys.exit(1)

import numpy as np

API_BASE_URL = os.environ.get('API_BASE_URL', 'http://localhost:8001')
API_KEY = os.environ.get('API_KEY', 'your-secret-key-here')
TEST_TIMEOUT = 120.0


class TTSASRIntegrationTest(unittest.TestCase):
    """TTS/ASR集成测试"""

    @classmethod
    def setUpClass(cls):
        cls.client = httpx.Client(timeout=TEST_TIMEOUT)
        cls.headers = {'X-API-Key': API_KEY}

        try:
            response = cls.client.get(f'{API_BASE_URL}/v1/tts-asr/status', headers=cls.headers)
            if response.status_code == 200:
                cls.service_available = True
                print(f"\n✓ 服务可用: {API_BASE_URL}")
            else:
                cls.service_available = False
                print(f"\n✗ 服务返回非200状态码: {response.status_code}")
        except Exception as e:  # noqa: ANN001
            cls.service_available = False
            print(f"\n✗ 无法连接到服务: {e}")

    @classmethod
    def tearDownClass(cls):
        cls.client.close()

    def setUp(self):
        if not self.service_available:
            self.skipTest("后端服务不可用")

    def test_01_config_endpoint(self):
        """测试配置端点"""
        response = self.client.get(
            f'{API_BASE_URL}/v1/tts-asr/config',
            headers=self.headers
        )

        self.assertEqual(response.status_code, 200)
        config = response.json()

        self.assertIn('device', config)
        self.assertIn('model', config)
        self.assertIn('status', config)

        model = config['model']
        status = config['status']
        self.assertIn('tts', model)
        self.assertIn('asr', model)

        print(f"\n配置信息:")
        print(f"  TTS模型: {model['tts']}")
        print(f"  ASR模型: {model.get('asr', 'N/A')}")
        print(f"  TTS已加载: {status['tts_loaded']}")
        print(f"  ASR已加载: {status['asr_loaded']}")

    def test_02_status_endpoint(self):
        """测试状态端点"""
        response = self.client.get(
            f'{API_BASE_URL}/v1/tts-asr/status',
            headers=self.headers
        )

        self.assertEqual(response.status_code, 200)
        status = response.json()

        self.assertIn('tts_loaded', status)
        self.assertIn('asr_loaded', status)
        self.assertIn('device', status)

        print(f"\n状态信息:")
        print(f"  TTS已加载: {status['tts_loaded']}")
        print(f"  ASR已加载: {status['asr_loaded']}")
        print(f"  设备: {status['device']}")

    def test_03_warmup_endpoint(self):
        """测试预热端点"""
        print("\n开始模型预热（可能需要几分钟）...")
        start_time = time.time()

        response = self.client.post(
            f'{API_BASE_URL}/v1/tts-asr/warmup',
            headers=self.headers,
        )

        elapsed = time.time() - start_time

        self.assertEqual(response.status_code, 200)
        result = response.json()

        self.assertIn('tts_warmup', result)
        self.assertIn('asr_warmup', result)

        print(f"\n预热完成 (耗时: {elapsed:.2f}秒):")
        print(f"  TTS预热: {'成功' if result['tts_warmup'] else '失败'}")
        print(f"  ASR预热: {'成功' if result.get('asr_warmup') else '失败/跳过'}")

        if not result['tts_warmup'] or not result.get('asr_warmup'):
            print("\n⚠ 警告: 预热失败可能是因为模型未下载")

    def test_04_tts_endpoint_basic(self):
        """测试TTS基本功能"""
        test_text = "这是一个测试"

        response = self.client.post(
            f'{API_BASE_URL}/v1/tts-asr/tts',
            headers=self.headers,
            json={'text': test_text}
        )

        if response.status_code == 500:
            error = response.json()
            print(f"\n⚠ TTS失败（可能是模型未加载）: {error.get('detail', 'Unknown error')}")
            self.skipTest("TTS模型未加载或不可用")

        self.assertEqual(response.status_code, 200)
        result = response.json()

        self.assertIn('audio_base64', result)
        self.assertIn('format', result)
        self.assertIn('duration_ms', result)

        audio_data = base64.b64decode(result['audio_base64'])
        self.assertGreater(len(audio_data), 0)

        print(f"\nTTS测试成功:")
        print(f"  输入文本: {test_text}")
        print(f"  音频大小: {len(audio_data)} bytes")

    def test_05_asr_endpoint_basic(self):
        """测试ASR基本功能"""
        sample_rate = 16000
        duration = 1.0
        samples = int(sample_rate * duration)

        silence = np.zeros(samples, dtype=np.int16)

        wav_buffer = io.BytesIO()
        with wave.open(wav_buffer, 'wb') as wf:  # noqa: SIM115
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(sample_rate)
            wf.writeframes(silence.tobytes())

        audio_bytes = wav_buffer.getvalue()
        audio_base64 = base64.b64encode(audio_bytes).decode()

        response = self.client.post(
            f'{API_BASE_URL}/v1/tts-asr/asr',
            headers=self.headers,
            json={
                'audio_base64': audio_base64,
                'language': 'zh-CN'
            }
        )

        if response.status_code in (500, 501):
            detail = response.json().get('detail', 'Unknown')
            print(f"\n⚠ ASR失败: {detail}")
            self.skipTest("ASR模型未加载或不可用")

        self.assertEqual(response.status_code, 200)
        result = response.json()

        self.assertIn('text', result)
        self.assertIn('language', result)

        print(f"\nASR测试成功:")
        print(f"  识别文本: '{result['text']}'")
        print(f"  语言: {result['language']}")

    def test_06_api_key_validation(self):
        """测试API密钥验证"""
        wrong_headers = {'X-API-Key': 'wrong-api-key'}

        response = self.client.get(
            f'{API_BASE_URL}/v1/tts-asr/status',
            headers=wrong_headers,
        )

        self.assertEqual(response.status_code, 403)


class PerformanceTest(unittest.TestCase):
    """性能测试"""

    @classmethod
    def setUpClass(cls):
        cls.client = httpx.Client(timeout=TEST_TIMEOUT)
        cls.headers = {'X-API-Key': API_KEY}

        try:
            response = cls.client.get(f'{API_BASE_URL}/v1/tts-asr/status', headers=cls.headers)
            cls.service_available = response.status_code == 200
        except Exception:  # noqa: ANN001, S110
            cls.service_available = False

    @classmethod
    def tearDownClass(cls):
        cls.client.close()

    def setUp(self):
        if not self.service_available:
            self.skipTest("后端服务不可用")

    def test_tts_latency(self):
        """测试TTS延迟"""
        latencies = []
        for i in range(3):
            start = time.time()
            response = self.client.post(
                f'{API_BASE_URL}/v1/tts-asr/tts',
                headers=self.headers,
                json={'text': '测试延迟'}
            )
            elapsed = time.time() - start

            if response.status_code == 200:
                latencies.append(elapsed)

        if latencies:
            print(f"\nTTS延迟测试:")
            print(f"  平均: {sum(latencies)/len(latencies):.3f}s")
            print(f"  最小: {min(latencies):.3f}s / 最大: {max(latencies):.3f}s")


def run_tests(test_type: Optional[str] = None) -> bool:
    """运行测试"""
    loader = unittest.TestLoader()
    suite = unittest.TestSuite()

    TEST_MAP = {
        'config': ('TTSASRIntegrationTest', 'test_01_config_endpoint'),
        'status': ('TTSASRIntegrationTest', 'test_02_status_endpoint'),
        'warmup': ('TTSASRIntegrationTest', 'test_03_warmup_endpoint'),
        'tts': ('TTSASRIntegrationTest', 'test_04_tts_endpoint_basic'),
        'asr': ('TTSASRIntegrationTest', 'test_05_asr_endpoint_basic'),
        'perf': ('PerformanceTest', None),
    }

    if test_type and test_type in TEST_MAP:
        cls_name, method = TEST_MAP[test_type]
        if method:
            suite.addTest(globals()[cls_name](method))
        else:
            suite.addTests(loader.loadTestsFromTestCase(globals()[cls_name]))
    elif test_type == 'api_key':
        suite.addTest(TTSASRIntegrationTest('test_06_api_key_validation'))
    else:
        suite.addTests(loader.loadTestsFromTestCase(TTSASRIntegrationTest))
        suite.addTests(loader.loadTestsFromTestCase(PerformanceTest))

    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(suite)
    return result.wasSuccessful()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='TTS/ASR 集成测试')
    parser.add_argument('--test', choices=['config', 'status', 'warmup', 'tts', 'asr', 'perf', 'api_key'])
    parser.add_argument('--url', default=API_BASE_URL)
    parser.add_argument('--key', default=API_KEY)

    args = parser.parse_args()
    API_BASE_URL = args.url
    API_KEY = args.key

    print("=" * 70)
    print("TTS/ASR 集成测试 (MLX/Qwen3-ASR)")
    print("=" * 70)

    success = run_tests(args.test)
    sys.exit(0 if success else 1)
