# -*- coding: utf-8 -*-
"""
AIfeng/2025-01-08 11:24:00
FunASR综合功能测试
测试内容：
1. 确认FunASR服务连接状态
2. 收音数据统计记录
3. 音频数据转发FunASR记录
4. 识别结果数据记录
"""

import os
import sys
import time
import json
import threading
import wave
import pyaudio
from datetime import datetime
from typing import Dict, List, Any

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from funasr_asr_sync import FunASRSync
from utils import config_util as cfg
from utils import util

class FunASRComprehensiveTest:
    """FunASR综合功能测试类"""
    
    def __init__(self):
        self.test_results = {
            'connection_test': {'status': 'pending', 'details': {}},
            'audio_stats': {'status': 'pending', 'details': {}},
            'audio_forwarding': {'status': 'pending', 'details': {}},
            'recognition_results': {'status': 'pending', 'details': {}}
        }
        
        # 音频统计数据
        self.audio_stats = {
            'total_frames': 0,
            'total_duration': 0.0,
            'audio_files_sent': 0,
            'bytes_processed': 0,
            'start_time': None,
            'end_time': None
        }
        
        # 转发记录
        self.forwarding_records = []
        
        # 识别结果记录
        self.recognition_records = []
        
        # 测试配置
        self.test_duration = 10  # 测试时长（秒）
        self.sample_rate = 16000
        self.chunk_size = 1024
        self.channels = 1
        self.format = pyaudio.paInt16
        
        # FunASR客户端
        self.asr_client = None
        
        # 录音相关
        self.audio = None
        self.stream = None
        self.is_recording = False
        
        print("FunASR综合测试初始化完成")
    
    def test_funasr_connection(self) -> bool:
        """测试1: 确认FunASR服务连接状态"""
        print("\n=== 测试1: FunASR服务连接状态 ===")
        
        try:
            # 创建FunASR客户端
            self.asr_client = FunASRSync("test_user")
            
            # 启动连接
            self.asr_client.start()
            
            # 等待连接建立
            connection_timeout = 10
            start_time = time.time()
            
            while time.time() - start_time < connection_timeout:
                if self.asr_client.is_connected():
                    self.test_results['connection_test'] = {
                        'status': 'success',
                        'details': {
                            'connected': True,
                            'connection_time': time.time() - start_time,
                            'server_url': f"ws://{cfg.local_asr_ip}:{cfg.local_asr_port}",
                            'timestamp': datetime.now().isoformat()
                        }
                    }
                    print(f"✓ FunASR服务连接成功")
                    print(f"  服务地址: ws://{cfg.local_asr_ip}:{cfg.local_asr_port}")
                    print(f"  连接耗时: {time.time() - start_time:.2f}秒")
                    return True
                time.sleep(0.5)
            
            # 连接超时
            self.test_results['connection_test'] = {
                'status': 'failed',
                'details': {
                    'connected': False,
                    'error': 'Connection timeout',
                    'timeout_duration': connection_timeout,
                    'server_url': f"ws://{cfg.local_asr_ip}:{cfg.local_asr_port}",
                    'timestamp': datetime.now().isoformat()
                }
            }
            print(f"✗ FunASR服务连接超时（{connection_timeout}秒）")
            return False
            
        except Exception as e:
            self.test_results['connection_test'] = {
                'status': 'failed',
                'details': {
                    'connected': False,
                    'error': str(e),
                    'error_type': type(e).__name__,
                    'timestamp': datetime.now().isoformat()
                }
            }
            print(f"✗ FunASR服务连接失败: {e}")
            return False
    
    def test_audio_recording_stats(self) -> bool:
        """测试2: 收音数据统计"""
        print("\n=== 测试2: 收音数据统计 ===")
        
        try:
            # 初始化音频
            self.audio = pyaudio.PyAudio()
            
            # 创建音频流
            self.stream = self.audio.open(
                format=self.format,
                channels=self.channels,
                rate=self.sample_rate,
                input=True,
                frames_per_buffer=self.chunk_size
            )
            
            print(f"开始录音测试，时长: {self.test_duration}秒")
            
            # 开始统计
            self.audio_stats['start_time'] = time.time()
            self.is_recording = True
            
            # 录音循环
            while time.time() - self.audio_stats['start_time'] < self.test_duration:
                try:
                    data = self.stream.read(self.chunk_size, exception_on_overflow=False)
                    
                    # 统计数据
                    self.audio_stats['total_frames'] += 1
                    self.audio_stats['bytes_processed'] += len(data)
                    
                    # 每秒输出一次统计
                    current_time = time.time()
                    if int(current_time - self.audio_stats['start_time']) % 2 == 0:
                        elapsed = current_time - self.audio_stats['start_time']
                        print(f"  录音进行中... {elapsed:.1f}s, 帧数: {self.audio_stats['total_frames']}, 数据量: {self.audio_stats['bytes_processed']} bytes")
                        time.sleep(0.1)  # 避免重复输出
                    
                except Exception as e:
                    print(f"录音数据读取错误: {e}")
                    break
            
            self.audio_stats['end_time'] = time.time()
            self.audio_stats['total_duration'] = self.audio_stats['end_time'] - self.audio_stats['start_time']
            
            # 停止录音
            self.is_recording = False
            self.stream.stop_stream()
            self.stream.close()
            self.audio.terminate()
            
            # 计算统计结果
            avg_frame_rate = self.audio_stats['total_frames'] / self.audio_stats['total_duration']
            data_rate_kbps = (self.audio_stats['bytes_processed'] * 8) / (self.audio_stats['total_duration'] * 1000)
            
            self.test_results['audio_stats'] = {
                'status': 'success',
                'details': {
                    'total_frames': self.audio_stats['total_frames'],
                    'total_duration': self.audio_stats['total_duration'],
                    'bytes_processed': self.audio_stats['bytes_processed'],
                    'avg_frame_rate': avg_frame_rate,
                    'data_rate_kbps': data_rate_kbps,
                    'sample_rate': self.sample_rate,
                    'channels': self.channels,
                    'timestamp': datetime.now().isoformat()
                }
            }
            
            print(f"✓ 收音数据统计完成")
            print(f"  总帧数: {self.audio_stats['total_frames']}")
            print(f"  总时长: {self.audio_stats['total_duration']:.2f}秒")
            print(f"  数据量: {self.audio_stats['bytes_processed']} bytes")
            print(f"  平均帧率: {avg_frame_rate:.1f} fps")
            print(f"  数据速率: {data_rate_kbps:.1f} kbps")
            
            return True
            
        except Exception as e:
            self.test_results['audio_stats'] = {
                'status': 'failed',
                'details': {
                    'error': str(e),
                    'error_type': type(e).__name__,
                    'timestamp': datetime.now().isoformat()
                }
            }
            print(f"✗ 收音数据统计失败: {e}")
            return False
    
    def test_audio_forwarding(self) -> bool:
        """测试3: 音频数据转发FunASR记录"""
        print("\n=== 测试3: 音频数据转发FunASR ===")
        
        if not self.asr_client or not self.asr_client.is_connected():
            print("✗ FunASR未连接，无法进行转发测试")
            self.test_results['audio_forwarding'] = {
                'status': 'failed',
                'details': {'error': 'FunASR not connected'}
            }
            return False
        
        try:
            # 创建测试音频文件
            test_audio_dir = "cache_data"
            os.makedirs(test_audio_dir, exist_ok=True)
            
            # 生成测试音频数据
            duration = 2.0  # 2秒测试音频
            frames = int(self.sample_rate * duration)
            
            # 创建简单的正弦波测试音频
            import numpy as np
            frequency = 440  # A4音符
            t = np.linspace(0, duration, frames, False)
            audio_data = np.sin(2 * np.pi * frequency * t) * 0.3
            audio_bytes = (audio_data * 32767).astype(np.int16).tobytes()
            
            # 保存为WAV文件
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f')[:-3]
            test_file = os.path.join(test_audio_dir, f'test_audio_{timestamp}.wav')
            
            with wave.open(test_file, 'wb') as wf:
                wf.setnchannels(self.channels)
                wf.setsampwidth(2)  # 16位
                wf.setframerate(self.sample_rate)
                wf.writeframes(audio_bytes)
            
            print(f"创建测试音频文件: {test_file}")
            
            # 记录转发信息
            forwarding_record = {
                'file_path': test_file,
                'file_size': os.path.getsize(test_file),
                'duration': duration,
                'sample_rate': self.sample_rate,
                'channels': self.channels,
                'send_time': time.time(),
                'timestamp': datetime.now().isoformat()
            }
            
            # 发送到FunASR
            print(f"发送音频文件到FunASR: {test_file}")
            self.asr_client.send_url(test_file)
            
            # 记录发送成功
            forwarding_record['sent_successfully'] = True
            self.forwarding_records.append(forwarding_record)
            self.audio_stats['audio_files_sent'] += 1
            
            self.test_results['audio_forwarding'] = {
                'status': 'success',
                'details': {
                    'files_sent': len(self.forwarding_records),
                    'last_file': test_file,
                    'last_file_size': forwarding_record['file_size'],
                    'forwarding_records': self.forwarding_records,
                    'timestamp': datetime.now().isoformat()
                }
            }
            
            print(f"✓ 音频数据转发成功")
            print(f"  文件: {test_file}")
            print(f"  大小: {forwarding_record['file_size']} bytes")
            print(f"  时长: {duration}秒")
            
            return True
            
        except Exception as e:
            self.test_results['audio_forwarding'] = {
                'status': 'failed',
                'details': {
                    'error': str(e),
                    'error_type': type(e).__name__,
                    'timestamp': datetime.now().isoformat()
                }
            }
            print(f"✗ 音频数据转发失败: {e}")
            return False
    
    def test_recognition_results(self) -> bool:
        """测试4: 识别结果数据记录"""
        print("\n=== 测试4: 识别结果数据记录 ===")
        
        if not self.asr_client:
            print("✗ FunASR客户端未初始化")
            self.test_results['recognition_results'] = {
                'status': 'failed',
                'details': {'error': 'ASR client not initialized'}
            }
            return False
        
        try:
            # 等待识别结果
            print("等待FunASR识别结果...")
            
            wait_timeout = 15  # 等待15秒
            start_wait = time.time()
            
            while time.time() - start_wait < wait_timeout:
                if self.asr_client.done and self.asr_client.finalResults:
                    # 记录识别结果
                    recognition_record = {
                        'result_text': self.asr_client.finalResults,
                        'receive_time': time.time(),
                        'processing_duration': time.time() - start_wait,
                        'timestamp': datetime.now().isoformat()
                    }
                    
                    self.recognition_records.append(recognition_record)
                    
                    self.test_results['recognition_results'] = {
                        'status': 'success',
                        'details': {
                            'results_received': len(self.recognition_records),
                            'last_result': recognition_record['result_text'],
                            'processing_duration': recognition_record['processing_duration'],
                            'recognition_records': self.recognition_records,
                            'timestamp': datetime.now().isoformat()
                        }
                    }
                    
                    print(f"✓ 识别结果接收成功")
                    print(f"  结果: {recognition_record['result_text']}")
                    print(f"  处理时长: {recognition_record['processing_duration']:.2f}秒")
                    
                    return True
                
                time.sleep(0.5)
            
            # 超时未收到结果
            self.test_results['recognition_results'] = {
                'status': 'failed',
                'details': {
                    'error': 'Recognition timeout',
                    'timeout_duration': wait_timeout,
                    'asr_done': self.asr_client.done,
                    'final_results': self.asr_client.finalResults,
                    'timestamp': datetime.now().isoformat()
                }
            }
            
            print(f"✗ 识别结果等待超时（{wait_timeout}秒）")
            print(f"  ASR状态: done={self.asr_client.done}, results='{self.asr_client.finalResults}'")
            
            return False
            
        except Exception as e:
            self.test_results['recognition_results'] = {
                'status': 'failed',
                'details': {
                    'error': str(e),
                    'error_type': type(e).__name__,
                    'timestamp': datetime.now().isoformat()
                }
            }
            print(f"✗ 识别结果记录失败: {e}")
            return False
    
    def run_comprehensive_test(self):
        """运行综合测试"""
        print("\n" + "="*50)
        print("FunASR综合功能测试开始")
        print("="*50)
        
        test_start_time = time.time()
        
        # 执行各项测试
        tests = [
            ("FunASR服务连接", self.test_funasr_connection),
            ("收音数据统计", self.test_audio_recording_stats),
            ("音频数据转发", self.test_audio_forwarding),
            ("识别结果记录", self.test_recognition_results)
        ]
        
        passed_tests = 0
        total_tests = len(tests)
        
        for test_name, test_func in tests:
            try:
                if test_func():
                    passed_tests += 1
            except Exception as e:
                print(f"✗ {test_name}测试异常: {e}")
        
        # 生成测试报告
        test_duration = time.time() - test_start_time
        
        print("\n" + "="*50)
        print("测试结果汇总")
        print("="*50)
        
        print(f"总测试数: {total_tests}")
        print(f"通过测试: {passed_tests}")
        print(f"失败测试: {total_tests - passed_tests}")
        print(f"测试耗时: {test_duration:.2f}秒")
        print(f"成功率: {(passed_tests/total_tests)*100:.1f}%")
        
        # 详细结果
        print("\n详细测试结果:")
        for i, (test_name, _) in enumerate(tests, 1):
            test_key = list(self.test_results.keys())[i-1]
            status = self.test_results[test_key]['status']
            status_symbol = "✓" if status == 'success' else "✗"
            print(f"  {i}. {test_name}: {status_symbol} {status}")
        
        # 保存测试报告
        self.save_test_report(test_duration, passed_tests, total_tests)
        
        # 清理资源
        self.cleanup()
        
        return passed_tests == total_tests
    
    def save_test_report(self, test_duration: float, passed_tests: int, total_tests: int):
        """保存测试报告"""
        try:
            report = {
                'test_summary': {
                    'total_tests': total_tests,
                    'passed_tests': passed_tests,
                    'failed_tests': total_tests - passed_tests,
                    'success_rate': (passed_tests/total_tests)*100,
                    'test_duration': test_duration,
                    'timestamp': datetime.now().isoformat()
                },
                'test_results': self.test_results,
                'audio_statistics': self.audio_stats,
                'forwarding_records': self.forwarding_records,
                'recognition_records': self.recognition_records
            }
            
            # 保存到文件
            report_file = f"test_funasr_comprehensive_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
            report_path = os.path.join("test", report_file)
            
            with open(report_path, 'w', encoding='utf-8') as f:
                json.dump(report, f, ensure_ascii=False, indent=2)
            
            print(f"\n测试报告已保存: {report_path}")
            
        except Exception as e:
            print(f"保存测试报告失败: {e}")
    
    def cleanup(self):
        """清理资源"""
        try:
            if self.asr_client:
                self.asr_client.end()
            
            if self.stream:
                self.stream.stop_stream()
                self.stream.close()
            
            if self.audio:
                self.audio.terminate()
                
            print("\n资源清理完成")
            
        except Exception as e:
            print(f"资源清理失败: {e}")

def main():
    """主函数"""
    try:
        # 检查依赖
        try:
            import numpy as np
        except ImportError:
            print("错误: 需要安装numpy库")
            print("请运行: pip install numpy")
            return False
        
        # 创建测试实例
        tester = FunASRComprehensiveTest()
        
        # 运行测试
        success = tester.run_comprehensive_test()
        
        if success:
            print("\n🎉 所有测试通过！")
            return True
        else:
            print("\n❌ 部分测试失败，请检查详细结果")
            return False
            
    except KeyboardInterrupt:
        print("\n测试被用户中断")
        return False
    except Exception as e:
        print(f"\n测试执行异常: {e}")
        return False

if __name__ == "__main__":
    main()