#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
macOS环境模拟测试工具
在非macOS环境下模拟Apple Silicon环境进行测试

运行方式:
    python backend/tests/simulate_macos.py --help
    python backend/tests/simulate_macos.py --device mps
    python backend/tests/simulate_macos.py --apple-silicon
    python backend/tests/simulate_macos.py --full-simulation
"""

import argparse
import importlib
import os
import platform
import sys
from unittest.mock import patch
import numpy as np


class MacOSSimulator:
    """macOS环境模拟器"""
    
    def __init__(self):
        self.original_platform_system = platform.system
        self.original_platform_machine = platform.machine
        self.patches = []
    
    def simulate_apple_silicon(self):
        """模拟Apple Silicon环境"""
        print("\n" + "="*70)
        print("模拟 Apple Silicon 环境")
        print("="*70)
        
        # 模拟Darwin系统和arm64架构
        self.patches.append(patch('platform.system', return_value='Darwin'))
        self.patches.append(patch('platform.machine', return_value='arm64'))
        
        for p in self.patches:
            p.start()
        
        print("✓ 平台: Darwin (macOS)")
        print("✓ 架构: arm64 (Apple Silicon)")
    
    def simulate_mps_device(self):
        """模拟MPS设备可用"""
        print("\n" + "="*70)
        print("模拟 MPS 设备")
        print("="*70)
        
        # 创建模拟的torch.backends.mps
        mock_mps = type('MockMPS', (), {
            'is_available': lambda: True,
            'is_built': lambda: True,
            'empty_cache': lambda: None
        })()
        
        mock_backends = type('MockBackends', (), {
            'mps': mock_mps
        })()
        
        # 模拟torch模块
        mock_torch = type('MockTorch', (), {
            'backends': mock_backends,
            'mps': mock_mps,
            'randn': lambda *args, **kwargs: np.random.randn(*args),
            'mm': lambda a, b: np.dot(a, b),
            'empty_cache': lambda: None
        })()
        
        self.patches.append(patch('torch', mock_torch))
        self.patches.append(patch('torch.backends.mps.is_available', return_value=True))
        self.patches.append(patch('torch.backends.mps.is_built', return_value=True))
        
        for p in self.patches[-3:]:
            p.start()
        
        print("✓ MPS 可用: True")
        print("✓ MPS 已编译: True")
    
    def simulate_cuda_device(self):
        """模拟CUDA设备可用"""
        print("\n" + "="*70)
        print("模拟 CUDA 设备")
        print("="*70)
        
        mock_cuda = type('MockCUDA', (), {
            'is_available': lambda: True,
            'device_count': lambda: 1,
            'get_device_properties': lambda n: type('Props', (), {'total_memory': 8*1024*1024*1024})(),
            'empty_cache': lambda: None
        })()
        
        self.patches.append(patch('torch.cuda', mock_cuda))
        self.patches.append(patch('torch.cuda.is_available', return_value=True))
        
        for p in self.patches[-2:]:
            p.start()
        
        print("✓ CUDA 可用: True")
        print("✓ GPU 数量: 1")
        print("✓ 显存: 8 GB")
    
    def cleanup(self):
        """清理所有补丁"""
        for p in self.patches:
            p.stop()
        self.patches.clear()
        print("\n✓ 已清理模拟环境")


def test_device_detection_on_apple_silicon():
    """测试Apple Silicon设备检测"""
    print("\n测试1: Apple Silicon 设备检测")
    print("-"*70)
    
    simulator = MacOSSimulator()
    try:
        simulator.simulate_apple_silicon()
        simulator.simulate_mps_device()
        
        # 设置环境变量
        os.environ['TTS_ASR_DEVICE'] = 'auto'
        os.environ['TTS_ASR_MODEL_SIZE'] = 'auto'
        
        # 重新导入模块以应用模拟
        if 'backend.tts_asr' in sys.modules:
            del sys.modules['backend.tts_asr']
        
        from backend.tts_asr import (
            _is_apple_silicon,
            _detect_device_capabilities,
            _get_recommended_model_size
        )
        
        # 测试Apple Silicon检测
        assert _is_apple_silicon(), "应该检测到Apple Silicon"
        print("✓ Apple Silicon 检测: 通过")
        
        # 测试设备能力检测
        caps = _detect_device_capabilities()
        print(f"✓ 设备: {caps.device}")
        print(f"✓ MPS 可用: {caps.mps_available}")
        print(f"✓ 推荐模型大小: {caps.recommended_model_size}")
        
        # 测试模型大小推荐
        recommended_size = _get_recommended_model_size()
        assert recommended_size in ['small', 'tiny', 'base'], \
            f"Apple Silicon应推荐小模型，但推荐了 {recommended_size}"
        print(f"✓ 推荐模型大小: {recommended_size}")
        
        print("\n✓ 测试通过")
        return True
        
    except Exception as e:
        print(f"\n✗ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        simulator.cleanup()


def test_memory_management():
    """测试内存管理"""
    print("\n测试2: 内存管理")
    print("-"*70)
    
    simulator = MacOSSimulator()
    try:
        simulator.simulate_apple_silicon()
        simulator.simulate_mps_device()
        
        # 模拟系统内存
        import psutil
        original_virtual_memory = psutil.virtual_memory
        
        def mock_virtual_memory():
            mock_mem = type('MockMemory', (), {
                'total': 16 * 1024 * 1024 * 1024  # 16GB
            })()
            return mock_mem
        
        self.patches.append(patch('psutil.virtual_memory', mock_virtual_memory))
        
        from backend.tts_asr import _get_system_memory_mb, TTS_ASR_MPS_MEMORY_LIMIT_MB
        
        mem_mb = _get_system_memory_mb()
        print(f"✓ 系统内存: {mem_mb} MB")
        
        # 计算预期的MPS内存限制（60%）
        expected_limit = int(mem_mb * 0.6)
        print(f"✓ 预期MPS限制: {expected_limit} MB (60%)")
        print(f"✓ 配置MPS限制: {TTS_ASR_MPS_MEMORY_LIMIT_MB} MB")
        
        print("\n✓ 测试通过")
        return True
        
    except Exception as e:
        print(f"\n✗ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        simulator.cleanup()


def test_model_size_selection():
    """测试模型大小选择"""
    print("\n测试3: 模型大小选择")
    print("-"*70)
    
    test_cases = [
        ('auto', 'Apple Silicon默认'),
        ('tiny', '最小模型'),
        ('small', '推荐模型'),
        ('medium', '中等模型'),
        ('large', '大模型'),
        ('turbo', 'turbo模型'),
    ]
    
    from backend.tts_asr import WHISPER_MODEL_SIZES, _get_recommended_model_size
    
    for size, desc in test_cases:
        os.environ['TTS_ASR_MODEL_SIZE'] = size
        
        # 重新加载模块
        if 'backend.tts_asr' in sys.modules:
            del sys.modules['backend.tts_asr']
        
        from backend.tts_asr import _get_recommended_model_size
        
        if size == 'auto':
            # 自动选择
            recommended = _get_recommended_model_size()
            print(f"✓ {desc}: {recommended}")
        else:
            # 显式选择
            os.environ['TTS_ASR_MODEL_SIZE'] = size
            result = _get_recommended_model_size()
            assert result == size, f"应该返回 {size}，但返回了 {result}"
            print(f"✓ {desc}: {size} -> {WHISPER_MODEL_SIZES[size]}")
    
    print("\n✓ 测试通过")
    return True


def test_audio_processing():
    """测试音频处理"""
    print("\n测试4: 音频处理")
    print("-"*70)
    
    from backend.tts_asr import (
        _validate_audio_data,
        _resample_audio_robust
    )
    
    # 测试音频验证
    test_cases = [
        (b'', False, "空数据"),
        (b'short', False, "太短"),
        (b'RIFF' + b'\x00' * 40, True, "有效WAV头"),
    ]
    
    for data, expected, desc in test_cases:
        result = _validate_audio_data(data)
        assert result == expected, f"{desc}: 预期 {expected}，得到 {result}"
        print(f"✓ 音频验证 ({desc}): {'通过' if result == expected else '失败'}")
    
    # 测试重采样
    audio_16k = np.sin(np.linspace(0, 2*np.pi, 16000)).astype(np.float32)
    
    # 16k -> 48k
    audio_48k = _resample_audio_robust(audio_16k, 16000, 48000)
    assert len(audio_48k) == 48000, f"48kHz音频长度错误: {len(audio_48k)}"
    print(f"✓ 重采样 (16k -> 48k): 长度 {len(audio_16k)} -> {len(audio_48k)}")
    
    # 48k -> 16k
    audio_back = _resample_audio_robust(audio_48k, 48000, 16000)
    assert len(audio_back) == 16000, f"16kHz音频长度错误: {len(audio_back)}"
    print(f"✓ 重采样 (48k -> 16k): 长度 {len(audio_48k)} -> {len(audio_back)}")
    
    print("\n✓ 测试通过")
    return True


def test_environment_variables():
    """测试环境变量"""
    print("\n测试5: 环境变量配置")
    print("-"*70)
    
    # 清理环境变量
    env_vars = [
        'TTS_ASR_DEVICE', 'TTS_ASR_MODEL_SIZE', 'TTS_ASR_QUANTIZE',
        'TTS_ASR_OFFLINE_MODE', 'TTS_ASR_WARMUP', 'TTS_ASR_WARMUP_TIMEOUT',
        'TTS_ASR_IDLE_TIMEOUT', 'TTS_ASR_MPS_MEMORY_LIMIT_MB'
    ]
    
    original_values = {}
    for var in env_vars:
        original_values[var] = os.environ.get(var)
        if var in os.environ:
            del os.environ[var]
    
    try:
        # 测试默认值
        from backend.tts_asr import (
            TTS_ASR_DEVICE, TTS_ASR_MODEL_SIZE, TTS_ASR_QUANTIZE,
            TTS_ASR_OFFLINE_MODE, TTS_ASR_WARMUP, TTS_ASR_WARMUP_TIMEOUT,
            TTS_ASR_IDLE_TIMEOUT, TTS_ASR_MPS_MEMORY_LIMIT_MB
        )
        
        defaults = {
            'TTS_ASR_DEVICE': 'auto',
            'TTS_ASR_MODEL_SIZE': 'auto',
            'TTS_ASR_QUANTIZE': False,
            'TTS_ASR_OFFLINE_MODE': False,
            'TTS_ASR_WARMUP': True,
            'TTS_ASR_WARMUP_TIMEOUT': 120,
            'TTS_ASR_IDLE_TIMEOUT': 0,
            'TTS_ASR_MPS_MEMORY_LIMIT_MB': 8192,
        }
        
        for var, expected in defaults.items():
            actual = locals()[var]
            assert actual == expected, f"{var}: 预期 {expected}，得到 {actual}"
            print(f"✓ {var} = {actual}")
        
        # 测试自定义值
        print("\n自定义配置测试:")
        os.environ['TTS_ASR_MODEL_SIZE'] = 'small'
        os.environ['TTS_ASR_QUANTIZE'] = 'true'
        os.environ['TTS_ASR_OFFLINE_MODE'] = 'true'
        os.environ['TTS_ASR_MPS_MEMORY_LIMIT_MB'] = '4096'
        
        # 重新加载
        if 'backend.tts_asr' in sys.modules:
            del sys.modules['backend.tts_asr']
        
        from backend.tts_asr import (
            TTS_ASR_MODEL_SIZE, TTS_ASR_QUANTIZE,
            TTS_ASR_OFFLINE_MODE, TTS_ASR_MPS_MEMORY_LIMIT_MB
        )
        
        assert TTS_ASR_MODEL_SIZE == 'small'
        assert TTS_ASR_QUANTIZE == True
        assert TTS_ASR_OFFLINE_MODE == True
        assert TTS_ASR_MPS_MEMORY_LIMIT_MB == 4096
        
        print(f"✓ TTS_ASR_MODEL_SIZE = {TTS_ASR_MODEL_SIZE}")
        print(f"✓ TTS_ASR_QUANTIZE = {TTS_ASR_QUANTIZE}")
        print(f"✓ TTS_ASR_OFFLINE_MODE = {TTS_ASR_OFFLINE_MODE}")
        print(f"✓ TTS_ASR_MPS_MEMORY_LIMIT_MB = {TTS_ASR_MPS_MEMORY_LIMIT_MB}")
        
        print("\n✓ 测试通过")
        return True
        
    finally:
        # 恢复原始值
        for var, value in original_values.items():
            if value is not None:
                os.environ[var] = value
            elif var in os.environ:
                del os.environ[var]


def run_full_simulation():
    """运行完整模拟测试"""
    print("\n" + "="*70)
    print("完整macOS环境模拟测试")
    print("="*70)
    
    results = []
    
    # 运行所有测试
    results.append(("设备检测", test_device_detection_on_apple_silicon()))
    results.append(("内存管理", test_memory_management()))
    results.append(("模型选择", test_model_size_selection()))
    results.append(("音频处理", test_audio_processing()))
    results.append(("环境变量", test_environment_variables()))
    
    # 汇总结果
    print("\n" + "="*70)
    print("测试结果汇总")
    print("="*70)
    
    for name, passed in results:
        status = "✓ 通过" if passed else "✗ 失败"
        print(f"{name}: {status}")
    
    total = len(results)
    passed = sum(1 for _, p in results if p)
    
    print("\n" + "-"*70)
    print(f"总计: {passed}/{total} 测试通过")
    print("="*70)
    
    return all(p for _, p in results)


def main():
    parser = argparse.ArgumentParser(
        description='macOS环境模拟测试工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 运行完整模拟测试
  python backend/tests/simulate_macos.py --full-simulation
  
  # 仅模拟Apple Silicon环境
  python backend/tests/simulate_macos.py --apple-silicon
  
  # 仅模拟MPS设备
  python backend/tests/simulate_macos.py --device mps
  
  # 仅模拟CUDA设备
  python backend/tests/simulate_macos.py --device cuda
        """
    )
    
    parser.add_argument(
        '--full-simulation',
        action='store_true',
        help='运行完整模拟测试'
    )
    
    parser.add_argument(
        '--apple-silicon',
        action='store_true',
        help='模拟Apple Silicon环境'
    )
    
    parser.add_argument(
        '--device',
        choices=['mps', 'cuda'],
        help='模拟特定设备'
    )
    
    parser.add_argument(
        '--test',
        choices=['device', 'memory', 'model', 'audio', 'env'],
        help='运行特定测试'
    )
    
    args = parser.parse_args()
    
    # 确保可以导入backend模块
    sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
    
    if args.full_simulation:
        success = run_full_simulation()
        sys.exit(0 if success else 1)
    
    if args.apple_silicon:
        simulator = MacOSSimulator()
        try:
            simulator.simulate_apple_silicon()
            simulator.simulate_mps_device()
            
            print("\n环境已模拟，按Ctrl+D退出")
            print("在Python环境中可以使用:")
            print("  from backend.tts_asr import _is_apple_silicon")
            print("  print(_is_apple_silicon())  # 应该返回 True")
            
            # 进入交互模式
            import code
            code.interact(local=locals())
        finally:
            simulator.cleanup()
    
    if args.device:
        simulator = MacOSSimulator()
        try:
            if args.device == 'mps':
                simulator.simulate_mps_device()
            elif args.device == 'cuda':
                simulator.simulate_cuda_device()
            
            print("\n设备已模拟")
            import code
            code.interact(local=locals())
        finally:
            simulator.cleanup()
    
    if args.test:
        test_func = {
            'device': test_device_detection_on_apple_silicon,
            'memory': test_memory_management,
            'model': test_model_size_selection,
            'audio': test_audio_processing,
            'env': test_environment_variables,
        }
        
        success = test_func[args.test]()
        sys.exit(0 if success else 1)
    
    # 默认运行完整测试
    if not any([args.full_simulation, args.apple_silicon, args.device, args.test]):
        parser.print_help()


if __name__ == '__main__':
    main()