# AIfeng/2025-07-07 15:25:48
# 识别结果追踪模块 - 流式识别结果的完整追踪与关联管理

import time
import uuid
import hashlib
import threading
from typing import List, Dict, Optional, Set, Tuple
from dataclasses import dataclass, field
from enum import Enum
import logging
from collections import defaultdict, deque

class ResultType(Enum):
    """识别结果类型"""
    PARTIAL = "partial"          # 部分结果
    REFINED = "refined"          # 精化结果
    FINAL = "final"              # 最终结果
    CORRECTED = "corrected"      # 修正结果

class ResultStatus(Enum):
    """结果状态"""
    ACTIVE = "active"            # 活跃状态
    SUPERSEDED = "superseded"    # 被替代
    EXPIRED = "expired"          # 已过期
    ARCHIVED = "archived"        # 已归档

@dataclass
class RecognitionSegmentID:
    """识别片段唯一标识"""
    session_id: str              # 会话ID
    segment_id: str              # 片段ID
    sequence_number: int         # 序列号
    parent_segment_id: Optional[str] = None  # 父片段ID
    timestamp: float = field(default_factory=time.time)
    
    def __post_init__(self):
        if not self.segment_id:
            self.segment_id = f"{self.session_id}_{self.sequence_number}_{int(self.timestamp * 1000)}"

@dataclass
class RecognitionResult:
    """增强的识别结果"""
    id: RecognitionSegmentID
    text: str
    confidence: float
    timestamp: float
    audio_duration: float
    result_type: ResultType
    stage: str                   # 识别阶段
    audio_segment_hash: str      # 音频片段哈希值
    
    # 关联信息
    predecessor_ids: List[str] = field(default_factory=list)
    successor_ids: List[str] = field(default_factory=list)
    
    # 状态信息
    status: ResultStatus = ResultStatus.ACTIVE
    superseded_by: Optional[str] = None
    superseded_at: Optional[float] = None
    
    # 质量指标
    accuracy_score: float = 0.0
    processing_time: float = 0.0
    
    # 元数据
    metadata: Dict = field(default_factory=dict)
    
    def to_dict(self) -> Dict:
        """转换为字典格式"""
        return {
            'id': {
                'session_id': self.id.session_id,
                'segment_id': self.id.segment_id,
                'sequence_number': self.id.sequence_number,
                'parent_segment_id': self.id.parent_segment_id,
                'timestamp': self.id.timestamp
            },
            'text': self.text,
            'confidence': self.confidence,
            'timestamp': self.timestamp,
            'audio_duration': self.audio_duration,
            'result_type': self.result_type.value,
            'stage': self.stage,
            'audio_segment_hash': self.audio_segment_hash,
            'predecessor_ids': self.predecessor_ids,
            'successor_ids': self.successor_ids,
            'status': self.status.value,
            'superseded_by': self.superseded_by,
            'superseded_at': self.superseded_at,
            'accuracy_score': self.accuracy_score,
            'processing_time': self.processing_time,
            'metadata': self.metadata
        }

class ResultRelationship:
    """结果关系管理"""
    
    def __init__(self):
        self.parent_child_map = defaultdict(set)  # 父子关系
        self.predecessor_successor_map = defaultdict(set)  # 前后继关系
        self.supersede_map = {}  # 替代关系
        
    def add_parent_child_relation(self, parent_id: str, child_id: str):
        """添加父子关系"""
        self.parent_child_map[parent_id].add(child_id)
    
    def add_predecessor_successor_relation(self, predecessor_id: str, successor_id: str):
        """添加前后继关系"""
        self.predecessor_successor_map[predecessor_id].add(successor_id)
    
    def add_supersede_relation(self, old_id: str, new_id: str):
        """添加替代关系"""
        self.supersede_map[old_id] = new_id
    
    def get_children(self, parent_id: str) -> Set[str]:
        """获取子结果"""
        return self.parent_child_map.get(parent_id, set())
    
    def get_successors(self, predecessor_id: str) -> Set[str]:
        """获取后继结果"""
        return self.predecessor_successor_map.get(predecessor_id, set())
    
    def get_superseding_result(self, old_id: str) -> Optional[str]:
        """获取替代结果"""
        return self.supersede_map.get(old_id)

class RecognitionResultTracker:
    """识别结果追踪器"""
    
    def __init__(self, config: Dict = None):
        self.config = config or self._get_default_config()
        
        # 结果存储
        self.result_graph = {}  # 结果关联图
        self.session_results = defaultdict(list)  # 按会话组织的结果
        self.active_segments = {}  # 活跃片段
        self.completed_segments = {}  # 完成片段
        
        # 关系管理
        self.relationships = ResultRelationship()
        
        # 序列号管理
        self.session_sequences = defaultdict(int)
        
        # 性能统计
        self.statistics = {
            'total_results': 0,
            'active_sessions': 0,
            'superseded_results': 0,
            'average_chain_length': 0.0
        }
        
        self.logger = logging.getLogger(__name__)
        self._lock = threading.RLock()
        
        # 回调函数管理
        self.result_callbacks = []  # 结果回调
        self.error_callbacks = []   # 错误回调
        
        # 启动清理任务
        self._start_cleanup_task()
    
    def _get_default_config(self) -> Dict:
        """获取默认配置"""
        return {
            'max_chain_length': 10,
            'cleanup_interval': 120.0,
            'max_session_age': 3600.0,  # 1小时
            'max_result_age': 300.0,    # 5分钟
            'enable_relationship_tracking': True,
            'enable_quality_tracking': True
        }
    
    def create_session(self, session_id: str = None) -> str:
        """创建新会话"""
        if not session_id:
            session_id = str(uuid.uuid4())
        
        with self._lock:
            if session_id not in self.session_sequences:
                self.session_sequences[session_id] = 0
                self.statistics['active_sessions'] += 1
                self.logger.info(f"创建新会话: {session_id}")
        
        return session_id
    
    def register_result_callback(self, callback):
        """注册结果回调函数"""
        self.result_callbacks.append(callback)
        self.logger.debug("注册结果回调函数")
    
    def register_error_callback(self, callback):
        """注册错误回调函数"""
        self.error_callbacks.append(callback)
        self.logger.debug("注册错误回调函数")
    
    def _trigger_result_callbacks(self, session_id: str, result: RecognitionResult):
        """触发结果回调"""
        for callback in self.result_callbacks:
            try:
                callback(session_id, result)
            except Exception as e:
                self.logger.error(f"结果回调执行失败: {e}")
    
    def _trigger_error_callbacks(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 add_recognition_result(self, session_id: str, text: str, confidence: float,
                             audio_data: bytes, result_type: ResultType, stage: str,
                             predecessor_ids: List[str] = None,
                             parent_segment_id: str = None,
                             metadata: Dict = None) -> str:
        """添加识别结果并建立关联"""
        
        with self._lock:
            # 生成序列号
            sequence_number = self.session_sequences[session_id]
            self.session_sequences[session_id] += 1
            
            # 创建结果ID
            result_id = RecognitionSegmentID(
                session_id=session_id,
                segment_id="",  # 将在__post_init__中生成
                sequence_number=sequence_number,
                parent_segment_id=parent_segment_id
            )
            
            # 计算音频哈希
            audio_hash = hashlib.md5(audio_data).hexdigest() if audio_data else ""
            
            # 创建识别结果
            result = RecognitionResult(
                id=result_id,
                text=text,
                confidence=confidence,
                timestamp=time.time(),
                audio_duration=len(audio_data) / 32000.0 if audio_data else 0.0,  # 假设16kHz, 2字节
                result_type=result_type,
                stage=stage,
                audio_segment_hash=audio_hash,
                predecessor_ids=predecessor_ids or [],
                metadata=metadata or {}
            )
            
            # 添加到结果图
            segment_id = result.id.segment_id
            self.result_graph[segment_id] = {
                'result': result,
                'predecessors': predecessor_ids or [],
                'successors': [],
                'created_at': time.time(),
                'last_accessed': time.time()
            }
            
            # 建立关联关系
            self._establish_relationships(result, predecessor_ids, parent_segment_id)
            
            # 添加到会话结果
            self.session_results[session_id].append(segment_id)
            
            # 更新统计信息
            self.statistics['total_results'] += 1
            
            self.logger.debug(f"添加识别结果: {segment_id}, 类型: {result_type.value}")
            
            return segment_id
    
    def _establish_relationships(self, result: RecognitionResult, 
                               predecessor_ids: List[str], 
                               parent_segment_id: str):
        """建立结果关联关系"""
        segment_id = result.id.segment_id
        
        # 建立前后继关系
        if predecessor_ids:
            for pred_id in predecessor_ids:
                if pred_id in self.result_graph:
                    # 更新前驱的后继列表
                    self.result_graph[pred_id]['successors'].append(segment_id)
                    
                    # 添加到关系管理器
                    self.relationships.add_predecessor_successor_relation(pred_id, segment_id)
                    
                    # 如果是最终结果，标记前驱为被替代
                    if result.result_type == ResultType.FINAL:
                        self._mark_superseded(pred_id, segment_id)
        
        # 建立父子关系
        if parent_segment_id and parent_segment_id in self.result_graph:
            self.relationships.add_parent_child_relation(parent_segment_id, segment_id)
    
    def _mark_superseded(self, old_result_id: str, new_result_id: str):
        """标记结果为被替代"""
        if old_result_id in self.result_graph:
            old_result = self.result_graph[old_result_id]['result']
            old_result.status = ResultStatus.SUPERSEDED
            old_result.superseded_by = new_result_id
            old_result.superseded_at = time.time()
            
            # 添加替代关系
            self.relationships.add_supersede_relation(old_result_id, new_result_id)
            
            self.statistics['superseded_results'] += 1
            
            self.logger.debug(f"结果 {old_result_id} 被 {new_result_id} 替代")
    
    def get_result_chain(self, segment_id: str) -> List[RecognitionResult]:
        """获取完整的识别链路"""
        with self._lock:
            if segment_id not in self.result_graph:
                return []
            
            # 更新访问时间
            self.result_graph[segment_id]['last_accessed'] = time.time()
            
            chain = []
            visited = set()
            
            # 向前追溯到起始结果
            self._trace_backwards(segment_id, chain, visited)
            
            # 向后追溯到最终结果
            self._trace_forwards(segment_id, chain, visited)
            
            # 按时间戳排序
            chain.sort(key=lambda r: r.timestamp)
            
            return chain
    
    def _trace_backwards(self, segment_id: str, chain: List[RecognitionResult], visited: Set[str]):
        """向前追溯结果链"""
        if segment_id in visited or segment_id not in self.result_graph:
            return
        
        visited.add(segment_id)
        result_info = self.result_graph[segment_id]
        result = result_info['result']
        
        # 添加当前结果
        if result not in chain:
            chain.append(result)
        
        # 递归追溯前驱
        for pred_id in result_info['predecessors']:
            self._trace_backwards(pred_id, chain, visited)
    
    def _trace_forwards(self, segment_id: str, chain: List[RecognitionResult], visited: Set[str]):
        """向后追溯结果链"""
        if segment_id in visited or segment_id not in self.result_graph:
            return
        
        visited.add(segment_id)
        result_info = self.result_graph[segment_id]
        
        # 递归追溯后继
        for succ_id in result_info['successors']:
            if succ_id not in visited and succ_id in self.result_graph:
                succ_result = self.result_graph[succ_id]['result']
                if succ_result not in chain:
                    chain.append(succ_result)
                self._trace_forwards(succ_id, chain, visited)
    
    def get_session_results(self, session_id: str, include_superseded: bool = False) -> List[RecognitionResult]:
        """获取会话的所有结果"""
        with self._lock:
            if session_id not in self.session_results:
                return []
            
            results = []
            for segment_id in self.session_results[session_id]:
                if segment_id in self.result_graph:
                    result = self.result_graph[segment_id]['result']
                    
                    if include_superseded or result.status != ResultStatus.SUPERSEDED:
                        results.append(result)
            
            # 按序列号排序
            results.sort(key=lambda r: r.id.sequence_number)
            
            return results
    
    def get_active_results(self, session_id: str) -> List[RecognitionResult]:
        """获取会话的活跃结果"""
        return [
            result for result in self.get_session_results(session_id)
            if result.status == ResultStatus.ACTIVE
        ]
    
    def get_final_results(self, session_id: str) -> List[RecognitionResult]:
        """获取会话的最终结果"""
        return [
            result for result in self.get_session_results(session_id)
            if result.result_type == ResultType.FINAL and result.status == ResultStatus.ACTIVE
        ]
    
    def update_result_quality(self, segment_id: str, accuracy_score: float, processing_time: float):
        """更新结果质量指标"""
        with self._lock:
            if segment_id in self.result_graph:
                result = self.result_graph[segment_id]['result']
                result.accuracy_score = accuracy_score
                result.processing_time = processing_time
                
                self.logger.debug(f"更新结果质量: {segment_id}, 准确率: {accuracy_score:.2f}")
    
    def complete_session(self, session_id: str):
        """完成会话"""
        with self._lock:
            if session_id in self.session_results:
                # 将活跃结果移动到完成状态
                for segment_id in self.session_results[session_id]:
                    if segment_id in self.result_graph:
                        result = self.result_graph[segment_id]['result']
                        if result.status == ResultStatus.ACTIVE:
                            result.status = ResultStatus.ARCHIVED
                
                self.statistics['active_sessions'] -= 1
                self.logger.info(f"会话已完成: {session_id}")
    
    def _start_cleanup_task(self):
        """启动清理任务"""
        def cleanup_worker():
            while True:
                try:
                    # 使用get方法获取配置，如果不存在则使用默认值
                    cleanup_interval = self.config.get('cleanup_interval', 120.0)
                    time.sleep(cleanup_interval)
                    self._cleanup_expired_results()
                except Exception as e:
                    self.logger.error(f"清理任务出错: {e}")
        
        cleanup_thread = threading.Thread(target=cleanup_worker, daemon=True)
        cleanup_thread.start()
        self.logger.info("清理任务已启动")
    
    def _cleanup_expired_results(self):
        """清理过期结果"""
        current_time = time.time()
        max_result_age = self.config.get('max_result_age', 300.0)
        max_session_age = self.config.get('max_session_age', 3600.0)
        
        with self._lock:
            expired_segments = []
            expired_sessions = []
            
            # 查找过期结果
            for segment_id, result_info in self.result_graph.items():
                if current_time - result_info['created_at'] > max_result_age:
                    result = result_info['result']
                    if result.status in [ResultStatus.SUPERSEDED, ResultStatus.ARCHIVED]:
                        expired_segments.append(segment_id)
            
            # 查找过期会话
            for session_id, result_ids in self.session_results.items():
                if result_ids:
                    # 检查会话中最新结果的时间
                    latest_time = max(
                        self.result_graph[rid]['created_at'] 
                        for rid in result_ids 
                        if rid in self.result_graph
                    )
                    
                    if current_time - latest_time > max_session_age:
                        expired_sessions.append(session_id)
            
            # 清理过期结果
            for segment_id in expired_segments:
                del self.result_graph[segment_id]
            
            # 清理过期会话
            for session_id in expired_sessions:
                del self.session_results[session_id]
                if session_id in self.session_sequences:
                    del self.session_sequences[session_id]
            
            if expired_segments or expired_sessions:
                self.logger.info(f"清理完成: {len(expired_segments)} 个结果, {len(expired_sessions)} 个会话")
    
    def get_statistics(self) -> Dict:
        """获取统计信息"""
        with self._lock:
            # 计算平均链长度
            if self.result_graph:
                total_chain_length = 0
                chain_count = 0
                
                for segment_id in self.result_graph:
                    chain = self.get_result_chain(segment_id)
                    if chain:
                        total_chain_length += len(chain)
                        chain_count += 1
                
                avg_chain_length = total_chain_length / chain_count if chain_count > 0 else 0
                self.statistics['average_chain_length'] = avg_chain_length
            
            return self.statistics.copy()
    
    def export_session_data(self, session_id: str) -> Dict:
        """导出会话数据"""
        with self._lock:
            results = self.get_session_results(session_id, include_superseded=True)
            
            return {
                'session_id': session_id,
                'total_results': len(results),
                'results': [result.to_dict() for result in results],
                'export_timestamp': time.time()
            }
    
    def get_performance_stats(self) -> Dict:
        """获取性能统计"""
        with self._lock:
            total_results = len(self.result_graph)
            total_sessions = len(self.session_results)
            
            return {
                'total_sessions': total_sessions,
                'total_results': total_results,
                'average_results_per_session': total_results / total_sessions if total_sessions > 0 else 0.0,
                'active_sessions': self.statistics['active_sessions']
            }
    
    def reset(self):
        """重置追踪器"""
        with self._lock:
            self.result_graph.clear()
            self.session_results.clear()
            self.active_segments.clear()
            self.completed_segments.clear()
            self.session_sequences.clear()
            self.relationships = ResultRelationship()
            
            # 重置统计信息
            self.statistics = {
                'total_results': 0,
                'active_sessions': 0,
                'superseded_results': 0,
                'average_chain_length': 0.0
            }
            
            self.logger.info("识别结果追踪器已重置")