# AIfeng/2025-07-07 15:25:48
# 流式语音识别优化集成管理器

import json
import time
import threading
import logging
from typing import Dict, List, Optional, Callable, Any
from pathlib import Path
from dataclasses import dataclass
from enum import Enum
import asyncio
from concurrent.futures import ThreadPoolExecutor

from .intelligent_segmentation import IntelligentSentenceSegmentation, SpeechSegment
from .adaptive_vad_chunking import AdaptiveVADChunking, ChunkStrategy, AudioChunk
from .recognition_result_tracker import RecognitionResultTracker, ResultType
from .streaming_display_manager import StreamingDisplayManager, UpdateType, DisplayPriority

class OptimizationMode(Enum):
    """优化模式"""
    SPEED_FIRST = "speed_first"          # 速度优先
    ACCURACY_FIRST = "accuracy_first"    # 精度优先
    BALANCED = "balanced"                # 平衡模式
    ADAPTIVE = "adaptive"                # 自适应模式

class ProcessingStage(Enum):
    """处理阶段"""
    AUDIO_INPUT = "audio_input"
    VAD_CHUNKING = "vad_chunking"
    SEGMENTATION = "segmentation"
    RECOGNITION = "recognition"
    RESULT_TRACKING = "result_tracking"
    DISPLAY_UPDATE = "display_update"

@dataclass
class ProcessingContext:
    """处理上下文"""
    session_id: str
    audio_data: bytes
    sample_rate: int
    timestamp: float
    metadata: Dict = None

@dataclass
class OptimizationMetrics:
    """优化指标"""
    total_latency_ms: float
    segmentation_latency_ms: float
    chunking_latency_ms: float
    tracking_latency_ms: float
    display_latency_ms: float
    accuracy_score: float
    confidence_score: float
    processing_efficiency: float

class OptimizationManager:
    """流式语音识别优化管理器"""
    
    def __init__(self, config_path: str = None):
        # 加载配置
        self.config = self._load_config(config_path)
        
        # 初始化各个优化模块
        self.segmentation_module = IntelligentSentenceSegmentation(
            self.config.get('intelligent_segmentation', {})
        )
        
        self.chunking_module = AdaptiveVADChunking(
            self.config.get('adaptive_vad_chunking', {})
        )
        
        self.tracking_module = RecognitionResultTracker(
            self.config.get('recognition_result_tracker', {})
        )
        
        self.display_module = StreamingDisplayManager(
            self.config.get('streaming_display_manager', {})
        )
        
        # 优化模式
        self.current_mode = OptimizationMode.BALANCED
        
        # 性能监控
        self.performance_metrics = {}
        self.processing_stats = {
            'total_sessions': 0,
            'active_sessions': 0,
            'total_audio_processed_seconds': 0.0,
            'average_latency_ms': 0.0,
            'average_accuracy': 0.0
        }
        
        # 回调函数
        self.result_callbacks = []  # 识别结果回调
        self.error_callbacks = []   # 错误处理回调
        self.metrics_callbacks = [] # 性能指标回调
        
        # 线程池
        self.executor = ThreadPoolExecutor(
            max_workers=self.config.get('integration', {}).get('performance_coordination', {}).get('max_workers', 8),
            thread_name_prefix='OptimizationManager'
        )
        
        # 事件总线（简化实现）
        self.event_handlers = {}
        
        self.logger = logging.getLogger(__name__)
        self._lock = threading.RLock()
        self._running = True
        
        # 注册模块间的回调
        self._setup_inter_module_communication()
        
        self.logger.info("流式语音识别优化管理器初始化完成")
    
    def _load_config(self, config_path: str = None) -> Dict:
        """加载配置文件"""
        if config_path is None:
            config_path = Path(__file__).parent / "optimization_config.json"
        
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
            
            # 转换配置项：将cleanup_interval_minutes转换为cleanup_interval（秒）
            if 'recognition_result_tracker' in config:
                tracker_config = config['recognition_result_tracker']
                if 'cleanup_interval_minutes' in tracker_config:
                    # 将分钟转换为秒
                    tracker_config['cleanup_interval'] = tracker_config['cleanup_interval_minutes'] * 60
                    # 保留原配置项以兼容
            
            return config
        except Exception as e:
            self.logger.warning(f"加载配置文件失败: {e}，使用默认配置")
            return self._get_default_config()
    
    def _get_default_config(self) -> Dict:
        """获取默认配置"""
        return {
            'intelligent_segmentation': {},
            'adaptive_vad_chunking': {},
            'recognition_result_tracker': {},
            'streaming_display_manager': {},
            'integration': {
                'performance_coordination': {
                    'max_workers': 8
                }
            }
        }
    
    def _setup_inter_module_communication(self):
        """设置模块间通信"""
        # 注册显示更新回调
        self.tracking_module.register_result_callback(self._on_tracking_result)
        self.display_module.register_error_callback(self._on_display_error)
        
        # 注册分片质量反馈
        self.chunking_module.register_quality_callback(self._on_chunk_quality_feedback)
    
    def set_optimization_mode(self, mode: OptimizationMode):
        """设置优化模式"""
        # 类型检查
        if not isinstance(mode, OptimizationMode):
            raise TypeError(f"mode必须是OptimizationMode枚举类型，当前类型: {type(mode)}")
            
        self.current_mode = mode
        
        # 根据模式调整各模块参数
        if mode == OptimizationMode.SPEED_FIRST:
            self._configure_for_speed()
        elif mode == OptimizationMode.ACCURACY_FIRST:
            self._configure_for_accuracy()
        elif mode == OptimizationMode.BALANCED:
            self._configure_for_balance()
        elif mode == OptimizationMode.ADAPTIVE:
            self._configure_for_adaptive()
        
        self.logger.info(f"优化模式已设置为: {mode.value}")
    
    def _configure_for_speed(self):
        """配置速度优先模式"""
        # 配置快速分片策略
        self.chunking_module.set_strategy(ChunkStrategy.FAST_RESPONSE)
        
        # 配置快速断句
        self.segmentation_module.update_config({
            'silence_thresholds': {
                'short_pause': 0.2,
                'medium_pause': 0.5,
                'long_pause': 1.0,
                'sentence_break': 1.5
            }
        })
        
        # 配置立即显示刷新
        self.display_module.config['refresh_strategies']['default_strategy'] = 'immediate'
    
    def _configure_for_accuracy(self):
        """配置精度优先模式"""
        # 配置高精度分片策略
        self.chunking_module.set_strategy(ChunkStrategy.HIGH_ACCURACY)
        
        # 配置精确断句
        self.segmentation_module.update_config({
            'semantic_analysis': {
                'enabled': True,
                'similarity_threshold': 0.8,
                'context_window': 8
            }
        })
        
        # 配置批量显示刷新
        self.display_module.config['refresh_strategies']['default_strategy'] = 'batch'
    
    def _configure_for_balance(self):
        """配置平衡模式"""
        # 配置平衡分片策略
        self.chunking_module.set_strategy(ChunkStrategy.BALANCED)
        
        # 配置防抖显示刷新
        self.display_module.config['refresh_strategies']['default_strategy'] = 'debounced'
    
    def _configure_for_adaptive(self):
        """配置自适应模式"""
        # 配置自适应分片策略
        self.chunking_module.set_strategy(ChunkStrategy.ADAPTIVE)
        
        # 配置自适应显示刷新
        self.display_module.config['refresh_strategies']['default_strategy'] = 'adaptive'
    
    def register_result_callback(self, callback: Callable[[str, str, float, bool], None]):
        """注册识别结果回调"""
        self.result_callbacks.append(callback)
    
    def register_error_callback(self, callback: Callable[[str, Exception], None]):
        """注册错误处理回调"""
        self.error_callbacks.append(callback)
    
    def register_metrics_callback(self, callback: Callable[[str, OptimizationMetrics], None]):
        """注册性能指标回调"""
        self.metrics_callbacks.append(callback)
    
    def create_session(self, session_id: str, config: Dict = None) -> bool:
        """创建处理会话"""
        try:
            # 在各个模块中创建会话
            self.segmentation_module.create_session(session_id)
            self.chunking_module.create_session(session_id)
            self.tracking_module.create_session(session_id)
            
            with self._lock:
                self.processing_stats['total_sessions'] += 1
                self.processing_stats['active_sessions'] += 1
            
            self.logger.info(f"会话创建成功: {session_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"创建会话失败: {e}")
            self._handle_error(session_id, e)
            return False
    
    def process_audio(self, session_id: str, audio_data: bytes, 
                     sample_rate: int, timestamp: float = None) -> bool:
        """处理音频数据"""
        if timestamp is None:
            timestamp = time.time()
        
        context = ProcessingContext(
            session_id=session_id,
            audio_data=audio_data,
            sample_rate=sample_rate,
            timestamp=timestamp
        )
        
        # 异步处理音频
        self.executor.submit(self._process_audio_async, context)
        return True
    
    def complete_session(self, session_id: str) -> bool:
        """完成处理会话"""
        try:
            # 完成各个模块的会话
            self.segmentation_module.complete_session(session_id)
            self.chunking_module.complete_session(session_id)
            self.tracking_module.complete_session(session_id)
            
            with self._lock:
                if self.processing_stats['active_sessions'] > 0:
                    self.processing_stats['active_sessions'] -= 1
            
            self.logger.info(f"会话完成: {session_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"完成会话失败: {e}")
            self._handle_error(session_id, e)
            return False
    
    def _process_audio_async(self, context: ProcessingContext):
        """异步处理音频数据"""
        start_time = time.time()
        metrics = OptimizationMetrics(
            total_latency_ms=0,
            segmentation_latency_ms=0,
            chunking_latency_ms=0,
            tracking_latency_ms=0,
            display_latency_ms=0,
            accuracy_score=0,
            confidence_score=0,
            processing_efficiency=0
        )
        
        try:
            # 1. VAD分片处理
            chunk_start = time.time()
            chunks = self.chunking_module.process_audio(
                context.session_id,
                context.audio_data,
                context.sample_rate
            )
            metrics.chunking_latency_ms = (time.time() - chunk_start) * 1000
            
            # 2. 智能断句处理
            seg_start = time.time()
            for chunk in chunks:
                if chunk.is_speech:
                    # 这里应该调用ASR服务获取识别结果
                    # 为了演示，我们模拟一个识别结果
                    mock_text = f"模拟识别文本_{chunk.chunk_id}"
                    mock_confidence = 0.85
                    
                    # 进行智能断句
                    text_context = {
                        'session_id': context.session_id,
                        'timestamp': chunk.timestamp,
                        'confidence': mock_confidence,
                        'silence_duration': 0.0  # 默认值
                    }
                    segment_result = self.segmentation_module.process_text(
                        mock_text,
                        text_context
                    )
                    
                    # 3. 结果追踪
                    track_start = time.time()
                    if segment_result.get('success', False):
                        result_id = self.tracking_module.add_recognition_result(
                            context.session_id,
                            segment_result['text'],
                            segment_result['confidence'],
                            context.audio_data,  # audio_data
                            ResultType.PARTIAL if not segment_result.get('is_complete', False) else ResultType.FINAL,  # result_type
                            'processing',  # stage
                            None,  # predecessor_ids
                            None,  # parent_segment_id
                            {'timestamp': chunk.timestamp, 'duration': chunk.duration}  # metadata
                        )
                        
                        # 4. 显示更新
                        display_start = time.time()
                        self.display_module.update_display(
                            context.session_id,
                            result_id,
                            segment_result['text'],
                            UpdateType.REPLACE_FINAL if segment_result.get('is_complete', False) else UpdateType.APPEND,
                            segment_result['confidence'],
                            segment_result.get('is_complete', False),
                            DisplayPriority.HIGH if segment_result.get('is_complete', False) else DisplayPriority.NORMAL
                        )
                        metrics.display_latency_ms += (time.time() - display_start) * 1000
                    
                    metrics.tracking_latency_ms += (time.time() - track_start) * 1000
            
            metrics.segmentation_latency_ms = (time.time() - seg_start) * 1000
            
            # 计算总延迟和效率
            metrics.total_latency_ms = (time.time() - start_time) * 1000
            # 防止除零错误
            if metrics.total_latency_ms > 0:
                metrics.processing_efficiency = len(context.audio_data) / metrics.total_latency_ms
            else:
                metrics.processing_efficiency = 0.0
                self.logger.warning(f"处理延迟为0，无法计算处理效率 [{context.session_id}]")
            
            # 更新性能统计
            self._update_performance_stats(context.session_id, metrics)
            
            # 触发指标回调
            self._trigger_metrics_callbacks(context.session_id, metrics)
            
        except Exception as e:
            self.logger.error(f"处理音频时出错: {e}")
            self._handle_error(context.session_id, e)
    
    def _on_tracking_result(self, session_id: str, result_id: str, text: str, 
                           confidence: float, is_final: bool):
        """处理追踪模块的结果回调"""
        # 触发结果回调
        for callback in self.result_callbacks:
            try:
                callback(session_id, text, confidence, is_final)
            except Exception as e:
                self.logger.error(f"结果回调执行出错: {e}")
    
    def _on_display_error(self, session_id: str, error: Exception):
        """处理显示模块的错误回调"""
        self.logger.error(f"显示模块错误 [{session_id}]: {error}")
        self._handle_error(session_id, error)
    
    def _on_chunk_quality_feedback(self, session_id: str, chunk_id: str, 
                                  quality_score: float, metrics: Dict):
        """处理分片质量反馈"""
        # 根据质量反馈调整策略
        if quality_score < 0.5:
            self.logger.warning(f"分片质量较低 [{session_id}:{chunk_id}]: {quality_score}")
            # 可以在这里实施自适应调整
    
    def _handle_error(self, session_id: str, error: Exception):
        """处理错误"""
        for callback in self.error_callbacks:
            try:
                callback(session_id, error)
            except Exception as e:
                self.logger.error(f"错误回调执行出错: {e}")
    
    def _update_performance_stats(self, session_id: str, metrics: OptimizationMetrics):
        """更新性能统计"""
        with self._lock:
            # 更新平均延迟
            current_avg = self.processing_stats['average_latency_ms']
            total_sessions = self.processing_stats['total_sessions']
            
            if total_sessions > 0:
                new_avg = (current_avg * (total_sessions - 1) + metrics.total_latency_ms) / total_sessions
                self.processing_stats['average_latency_ms'] = new_avg
            
            # 存储会话指标
            self.performance_metrics[session_id] = metrics
    
    def _trigger_metrics_callbacks(self, session_id: str, metrics: OptimizationMetrics):
        """触发性能指标回调"""
        for callback in self.metrics_callbacks:
            try:
                # 将session_id包含在metrics字典中传递给回调
                metrics_dict = {
                    'session_id': session_id,
                    'total_latency_ms': metrics.total_latency_ms,
                    'chunking_latency_ms': metrics.chunking_latency_ms,
                    'segmentation_latency_ms': metrics.segmentation_latency_ms,
                    'tracking_latency_ms': metrics.tracking_latency_ms,
                    'display_latency_ms': metrics.display_latency_ms,
                    'processing_efficiency': metrics.processing_efficiency,
                    'accuracy_score': getattr(metrics, 'accuracy_score', 0.0)
                }
                callback(metrics_dict)
            except Exception as e:
                self.logger.error(f"指标回调执行出错: {e}")
    
    def complete_session(self, session_id: str) -> bool:
        """完成处理会话"""
        try:
            # 完成各个模块的会话
            self.segmentation_module.complete_session(session_id)
            self.chunking_module.complete_session(session_id)
            self.tracking_module.complete_session(session_id)
            
            with self._lock:
                self.processing_stats['active_sessions'] -= 1
                if session_id in self.performance_metrics:
                    del self.performance_metrics[session_id]
            
            self.logger.info(f"会话完成: {session_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"完成会话失败: {e}")
            self._handle_error(session_id, e)
            return False
    
    def get_session_results(self, session_id: str) -> List[Dict]:
        """获取会话的所有结果"""
        try:
            # 从追踪模块获取结果
            results = self.tracking_module.get_session_results(session_id)
            
            # 从显示模块获取显示信息
            display_segments = self.display_module.get_session_display(session_id)
            
            # 合并结果
            combined_results = []
            for result in results:
                result_dict = {
                    'result_id': result.result_id,
                    'text': result.text,
                    'confidence': result.confidence,
                    'is_final': result.is_final,
                    'timestamp': result.timestamp,
                    'result_type': result.result_type.value if hasattr(result.result_type, 'value') else str(result.result_type)
                }
                combined_results.append(result_dict)
            
            return combined_results
            
        except Exception as e:
            self.logger.error(f"获取会话结果失败: {e}")
            return []
    
    def get_performance_stats(self) -> Dict:
        """获取性能统计"""
        with self._lock:
            stats = self.processing_stats.copy()
            
            # 添加各模块的性能统计
            stats['segmentation_stats'] = self.segmentation_module.get_performance_stats()
            stats['chunking_stats'] = self.chunking_module.get_performance_stats()
            stats['tracking_stats'] = self.tracking_module.get_performance_stats()
            stats['display_stats'] = self.display_module.get_performance_stats()
            
            return stats
    
    def get_optimization_metrics(self, session_id: str = None) -> Dict:
        """获取优化指标"""
        if session_id:
            return self.performance_metrics.get(session_id, {})
        else:
            return self.performance_metrics.copy()
    
    def shutdown(self):
        """关闭优化管理器"""
        self._running = False
        
        # 关闭各个模块
        self.segmentation_module.shutdown()
        self.chunking_module.shutdown()
        self.tracking_module.shutdown()
        self.display_module.shutdown()
        
        # 关闭线程池
        self.executor.shutdown(wait=True)
        
        self.logger.info("流式语音识别优化管理器已关闭")