"""
协议处理器模块
负责处理各种网络协议和消息格式
"""

import time
import json
import struct
from typing import Dict, Any, List, Optional
import hashlib
import zlib

class ProtocolHandler:
    """
    协议处理器
    负责处理各种网络协议和消息格式
    """
    
    def __init__(self, plugin):
        """
        初始化协议处理器
        
        Args:
            plugin: 网络同步插件实例
        """
        self.plugin = plugin
        self.enabled = False
        self.initialized = False
        
        # 协议配置
        self.protocol_config = {
            'version': '1.0',
            'enable_compression': True,
            'compression_level': 6,
            'enable_checksum': True,
            'max_message_size': 65536,  # 64KB
            'enable_fragmentation': True,
            'fragment_size': 1024,
            'enable_encryption': False,
            'supported_protocols': ['sync_v1', 'sync_v2']
        }
        
        # 消息类型定义
        self.message_types = {
            'CONNECT': 1,
            'CONNECT_ACK': 2,
            'DISCONNECT': 3,
            'HEARTBEAT': 4,
            'SYNC_DATA': 5,
            'SYNC_REQUEST': 6,
            'SYNC_RESPONSE': 7,
            'RPC_CALL': 8,
            'RPC_RESPONSE': 9,
            'EVENT': 10,
            'CHAT': 11,
            'ERROR': 12
        }
        
        # 消息处理器
        self.message_handlers = {}
        
        # 协议状态
        self.protocol_state = {
            'current_protocol': 'sync_v1',
            'bytes_processed': 0,
            'messages_processed': 0,
            'compression_savings': 0,
            'fragmented_messages': 0
        }
        
        # 消息片段缓存
        self.message_fragments = {}
        
        # 协议统计
        self.protocol_stats = {
            'messages_sent': 0,
            'messages_received': 0,
            'bytes_sent': 0,
            'bytes_received': 0,
            'compressed_messages': 0,
            'checksum_errors': 0,
            'protocol_errors': 0,
            'fragmented_packets': 0
        }
        
        # 回调函数
        self.protocol_callbacks = {
            'message_received': [],
            'message_sent': [],
            'protocol_error': [],
            'message_fragmented': []
        }
        
        # 时间戳记录
        self.last_protocol_update = 0.0
        self.last_compression_savings = 0.0
        
        print("✓ 协议处理器已创建")
        
    def initialize(self) -> bool:
        """
        初始化协议处理器
        
        Returns:
            是否初始化成功
        """
        try:
            # 注册默认消息处理器
            self._register_default_handlers()
            
            self.initialized = True
            print("✓ 协议处理器初始化完成")
            return True
            
        except Exception as e:
            print(f"✗ 协议处理器初始化失败: {e}")
            import traceback
            traceback.print_exc()
            return False
            
    def enable(self) -> bool:
        """
        启用协议处理器
        
        Returns:
            是否启用成功
        """
        try:
            if not self.initialized:
                print("✗ 协议处理器未初始化")
                return False
                
            self.enabled = True
            print("✓ 协议处理器已启用")
            return True
            
        except Exception as e:
            print(f"✗ 协议处理器启用失败: {e}")
            import traceback
            traceback.print_exc()
            return False
            
    def disable(self):
        """禁用协议处理器"""
        try:
            self.enabled = False
            print("✓ 协议处理器已禁用")
            
        except Exception as e:
            print(f"✗ 协议处理器禁用失败: {e}")
            import traceback
            traceback.print_exc()
            
    def finalize(self):
        """清理协议处理器资源"""
        try:
            self.disable()
            self.message_handlers.clear()
            self.message_fragments.clear()
            self.protocol_callbacks.clear()
            self.initialized = False
            print("✓ 协议处理器资源已清理")
            
        except Exception as e:
            print(f"✗ 协议处理器资源清理失败: {e}")
            import traceback
            traceback.print_exc()
            
    def update(self, dt: float):
        """
        更新协议处理器状态
        
        Args:
            dt: 时间增量
        """
        try:
            if not self.enabled:
                return
                
            current_time = time.time()
            self.last_protocol_update = current_time
            
            # 清理过期的消息片段
            self._cleanup_fragments(current_time)
            
        except Exception as e:
            print(f"✗ 协议处理器更新失败: {e}")
            import traceback
            traceback.print_exc()
            
    def _register_default_handlers(self):
        """注册默认消息处理器"""
        try:
            self.message_handlers = {
                self.message_types['CONNECT']: self._handle_connect,
                self.message_types['CONNECT_ACK']: self._handle_connect_ack,
                self.message_types['DISCONNECT']: self._handle_disconnect,
                self.message_types['HEARTBEAT']: self._handle_heartbeat,
                self.message_types['SYNC_DATA']: self._handle_sync_data,
                self.message_types['SYNC_REQUEST']: self._handle_sync_request,
                self.message_types['SYNC_RESPONSE']: self._handle_sync_response,
                self.message_types['RPC_CALL']: self._handle_rpc_call,
                self.message_types['RPC_RESPONSE']: self._handle_rpc_response,
                self.message_types['EVENT']: self._handle_event,
                self.message_types['CHAT']: self._handle_chat,
                self.message_types['ERROR']: self._handle_error
            }
            
            print("✓ 默认消息处理器已注册")
        except Exception as e:
            print(f"✗ 默认消息处理器注册失败: {e}")
            
    def _cleanup_fragments(self, current_time: float):
        """
        清理过期的消息片段
        
        Args:
            current_time: 当前时间
        """
        try:
            expired_fragments = [
                msg_id for msg_id, fragment_data in self.message_fragments.items()
                if current_time - fragment_data.get('timestamp', 0) > 30.0  # 30秒超时
            ]
            
            for msg_id in expired_fragments:
                del self.message_fragments[msg_id]
                
        except Exception as e:
            print(f"✗ 消息片段清理失败: {e}")
            
    def encode_message(self, message_type: str, data: Dict[str, Any], 
                      compress: bool = None, add_checksum: bool = None) -> Optional[bytes]:
        """
        编码消息
        
        Args:
            message_type: 消息类型
            data: 消息数据
            compress: 是否压缩
            add_checksum: 是否添加校验和
            
        Returns:
            编码后的消息数据或None
        """
        try:
            # 获取消息类型ID
            type_id = self.message_types.get(message_type.upper())
            if type_id is None:
                print(f"✗ 未知的消息类型: {message_type}")
                return None
                
            # 序列化数据
            data_json = json.dumps(data, ensure_ascii=False)
            data_bytes = data_json.encode('utf-8')
            original_size = len(data_bytes)
            
            # 压缩数据
            compressed = False
            if compress is None:
                compress = self.protocol_config['enable_compression']
                
            if compress and original_size > 256:  # 只压缩大于256字节的数据
                compressed_data = zlib.compress(data_bytes, self.protocol_config['compression_level'])
                if len(compressed_data) < original_size:
                    data_bytes = compressed_data
                    compressed = True
                    self.protocol_stats['compressed_messages'] += 1
                    self.protocol_state['compression_savings'] += original_size - len(compressed_data)
                    
            # 添加校验和
            checksum = b''
            if add_checksum is None:
                add_checksum = self.protocol_config['enable_checksum']
                
            if add_checksum:
                checksum = hashlib.md5(data_bytes).digest()
                
            # 构造消息头
            timestamp = int(time.time() * 1000) & 0xFFFFFFFF  # 32位时间戳
            header = struct.pack('!BBI', type_id, compressed, timestamp)
            
            # 构造完整消息
            message = header + checksum + data_bytes
            
            # 更新统计信息
            self.protocol_stats['messages_sent'] += 1
            self.protocol_stats['bytes_sent'] += len(message)
            self.protocol_state['bytes_processed'] += len(message)
            self.protocol_state['messages_processed'] += 1
            
            return message
            
        except Exception as e:
            print(f"✗ 消息编码失败: {e}")
            self.protocol_stats['protocol_errors'] += 1
            return None
            
    def decode_message(self, message_data: bytes) -> Optional[Dict[str, Any]]:
        """
        解码消息
        
        Args:
            message_data: 消息数据
            
        Returns:
            解码后的消息字典或None
        """
        try:
            if len(message_data) < 6:  # 最小消息长度(1字节类型+1字节压缩标志+4字节时间戳)
                print("✗ 消息数据太短")
                self.protocol_stats['protocol_errors'] += 1
                return None
                
            # 解析消息头
            header_size = 6
            type_id, compressed_flag, timestamp = struct.unpack('!BBI', message_data[:header_size])
            
            # 获取消息类型名称
            message_type = None
            for name, id_val in self.message_types.items():
                if id_val == type_id:
                    message_type = name
                    break
                    
            if message_type is None:
                print(f"✗ 未知的消息类型ID: {type_id}")
                self.protocol_stats['protocol_errors'] += 1
                return None
                
            # 计算校验和偏移
            checksum_offset = header_size
            if self.protocol_config['enable_checksum']:
                checksum_offset += 16  # MD5校验和长度
                
            # 获取数据部分
            data_bytes = message_data[checksum_offset:]
            
            # 验证校验和
            if self.protocol_config['enable_checksum'] and len(message_data) >= checksum_offset:
                expected_checksum = message_data[header_size:checksum_offset]
                actual_checksum = hashlib.md5(data_bytes).digest()
                if expected_checksum != actual_checksum:
                    print("✗ 消息校验和错误")
                    self.protocol_stats['checksum_errors'] += 1
                    return None
                    
            # 解压缩数据
            if compressed_flag:
                try:
                    data_bytes = zlib.decompress(data_bytes)
                except Exception as e:
                    print(f"✗ 消息解压缩失败: {e}")
                    self.protocol_stats['protocol_errors'] += 1
                    return None
                    
            # 反序列化数据
            data_json = data_bytes.decode('utf-8')
            data = json.loads(data_json)
            
            # 构造解码结果
            result = {
                'type': message_type,
                'type_id': type_id,
                'compressed': bool(compressed_flag),
                'timestamp': timestamp,
                'data': data
            }
            
            # 更新统计信息
            self.protocol_stats['messages_received'] += 1
            self.protocol_stats['bytes_received'] += len(message_data)
            self.protocol_state['bytes_processed'] += len(message_data)
            self.protocol_state['messages_processed'] += 1
            
            # 触发消息接收回调
            self._trigger_protocol_callback('message_received', result)
            
            return result
            
        except Exception as e:
            print(f"✗ 消息解码失败: {e}")
            self.protocol_stats['protocol_errors'] += 1
            return None
            
    def fragment_message(self, message_data: bytes, fragment_size: int = None) -> List[bytes]:
        """
        分片消息
        
        Args:
            message_data: 消息数据
            fragment_size: 分片大小
            
        Returns:
            分片后的消息列表
        """
        try:
            if fragment_size is None:
                fragment_size = self.protocol_config['fragment_size']
                
            if len(message_data) <= fragment_size:
                return [message_data]
                
            fragments = []
            message_id = hashlib.md5(message_data).hexdigest()[:8]
            total_fragments = (len(message_data) + fragment_size - 1) // fragment_size
            
            for i in range(total_fragments):
                start_pos = i * fragment_size
                end_pos = min((i + 1) * fragment_size, len(message_data))
                fragment_data = message_data[start_pos:end_pos]
                
                # 构造分片头
                fragment_header = struct.pack('!8sHH', 
                                            message_id.encode('ascii'), 
                                            i, 
                                            total_fragments)
                
                fragments.append(fragment_header + fragment_data)
                
            self.protocol_stats['fragmented_packets'] += len(fragments)
            self.protocol_state['fragmented_messages'] += 1
            
            # 触发消息分片回调
            self._trigger_protocol_callback('message_fragmented', {
                'message_id': message_id,
                'total_fragments': total_fragments,
                'fragment_size': fragment_size
            })
            
            return fragments
            
        except Exception as e:
            print(f"✗ 消息分片失败: {e}")
            self.protocol_stats['protocol_errors'] += 1
            return [message_data]
            
    def reassemble_message(self, fragment_data: bytes) -> Optional[bytes]:
        """
        重组消息
        
        Args:
            fragment_data: 分片数据
            
        Returns:
            重组后的消息数据或None
        """
        try:
            if len(fragment_data) < 12:  # 分片头最小长度
                return None
                
            # 解析分片头
            message_id_bytes, fragment_index, total_fragments = struct.unpack('!8sHH', fragment_data[:12])
            message_id = message_id_bytes.decode('ascii').rstrip('\x00')
            fragment_payload = fragment_data[12:]
            
            # 初始化分片缓存
            if message_id not in self.message_fragments:
                self.message_fragments[message_id] = {
                    'fragments': {},
                    'total_fragments': total_fragments,
                    'timestamp': time.time()
                }
                
            # 存储分片
            self.message_fragments[message_id]['fragments'][fragment_index] = fragment_payload
            
            # 检查是否所有分片都已到达
            fragment_cache = self.message_fragments[message_id]
            if len(fragment_cache['fragments']) == fragment_cache['total_fragments']:
                # 重组消息
                message_parts = []
                for i in range(fragment_cache['total_fragments']):
                    if i in fragment_cache['fragments']:
                        message_parts.append(fragment_cache['fragments'][i])
                    else:
                        # 缺少分片
                        return None
                        
                # 清理缓存
                del self.message_fragments[message_id]
                
                # 返回重组后的消息
                return b''.join(message_parts)
                
            return None
            
        except Exception as e:
            print(f"✗ 消息重组失败: {e}")
            self.protocol_stats['protocol_errors'] += 1
            return None
            
    def send_message(self, message_type: str, data: Dict[str, Any], 
                    target_client_id: str = None) -> bool:
        """
        发送消息
        
        Args:
            message_type: 消息类型
            data: 消息数据
            target_client_id: 目标客户端ID
            
        Returns:
            是否发送成功
        """
        try:
            if not self.enabled:
                print("✗ 协议处理器未启用")
                return False
                
            # 编码消息
            message_bytes = self.encode_message(message_type, data)
            if message_bytes is None:
                print("✗ 消息编码失败")
                return False
                
            # 检查是否需要分片
            if (self.protocol_config['enable_fragmentation'] and 
                len(message_bytes) > self.protocol_config['fragment_size']):
                fragments = self.fragment_message(message_bytes)
                success = True
                
                # 发送所有分片
                for fragment in fragments:
                    if not self._send_raw_message(fragment, target_client_id):
                        success = False
                        
                if success:
                    # 触发消息发送回调
                    self._trigger_protocol_callback('message_sent', {
                        'type': message_type,
                        'data': data,
                        'fragmented': True,
                        'fragment_count': len(fragments)
                    })
                    
                return success
            else:
                # 发送完整消息
                if self._send_raw_message(message_bytes, target_client_id):
                    # 触发消息发送回调
                    self._trigger_protocol_callback('message_sent', {
                        'type': message_type,
                        'data': data,
                        'fragmented': False
                    })
                    return True
                else:
                    return False
                    
        except Exception as e:
            print(f"✗ 消息发送失败: {e}")
            self.protocol_stats['protocol_errors'] += 1
            return False
            
    def _send_raw_message(self, message_bytes: bytes, target_client_id: str = None) -> bool:
        """
        发送原始消息数据
        
        Args:
            message_bytes: 消息字节数据
            target_client_id: 目标客户端ID
            
        Returns:
            是否发送成功
        """
        try:
            # 通过网络管理器发送消息
            if self.plugin.network_manager:
                # 这里应该调用网络管理器的发送方法
                # 由于是模拟实现，我们只返回True
                return True
            else:
                print("✗ 网络管理器未初始化")
                return False
                
        except Exception as e:
            print(f"✗ 原始消息发送失败: {e}")
            return False
            
    def receive_message(self, message_bytes: bytes) -> bool:
        """
        接收消息
        
        Args:
            message_bytes: 消息字节数据
            
        Returns:
            是否处理成功
        """
        try:
            if not self.enabled:
                print("✗ 协议处理器未启用")
                return False
                
            # 检查是否为分片消息
            if len(message_bytes) >= 12:
                # 尝试解析分片头
                try:
                    message_id_bytes, fragment_index, total_fragments = struct.unpack('!8sHH', message_bytes[:12])
                    # 这是一个分片消息
                    reassembled_message = self.reassemble_message(message_bytes)
                    if reassembled_message is not None:
                        # 分片重组完成，处理完整消息
                        return self.receive_message(reassembled_message)
                    else:
                        # 等待更多分片
                        return True
                except struct.error:
                    # 不是分片消息，继续处理
                    pass
                    
            # 解码消息
            decoded_message = self.decode_message(message_bytes)
            if decoded_message is None:
                print("✗ 消息解码失败")
                return False
                
            # 处理消息
            return self._process_message(decoded_message)
            
        except Exception as e:
            print(f"✗ 消息接收失败: {e}")
            self.protocol_stats['protocol_errors'] += 1
            return False
            
    def _process_message(self, message: Dict[str, Any]) -> bool:
        """
        处理消息
        
        Args:
            message: 解码后的消息
            
        Returns:
            是否处理成功
        """
        try:
            message_type_id = message.get('type_id')
            message_data = message.get('data', {})
            
            # 查找对应的消息处理器
            if message_type_id in self.message_handlers:
                handler = self.message_handlers[message_type_id]
                return handler(message_data)
            else:
                print(f"✗ 未找到消息处理器: {message_type_id}")
                return False
                
        except Exception as e:
            print(f"✗ 消息处理失败: {e}")
            return False
            
    # 默认消息处理器
    def _handle_connect(self, data: Dict[str, Any]) -> bool:
        """处理连接消息"""
        try:
            print(f"✓ 处理连接消息: {data}")
            return True
        except Exception as e:
            print(f"✗ 连接消息处理失败: {e}")
            return False
            
    def _handle_connect_ack(self, data: Dict[str, Any]) -> bool:
        """处理连接确认消息"""
        try:
            print(f"✓ 处理连接确认消息: {data}")
            return True
        except Exception as e:
            print(f"✗ 连接确认消息处理失败: {e}")
            return False
            
    def _handle_disconnect(self, data: Dict[str, Any]) -> bool:
        """处理断开连接消息"""
        try:
            print(f"✓ 处理断开连接消息: {data}")
            return True
        except Exception as e:
            print(f"✗ 断开连接消息处理失败: {e}")
            return False
            
    def _handle_heartbeat(self, data: Dict[str, Any]) -> bool:
        """处理心跳消息"""
        try:
            print(f"✓ 处理心跳消息: {data}")
            return True
        except Exception as e:
            print(f"✗ 心跳消息处理失败: {e}")
            return False
            
    def _handle_sync_data(self, data: Dict[str, Any]) -> bool:
        """处理同步数据消息"""
        try:
            print(f"✓ 处理同步数据消息: {len(data)}个对象")
            
            # 传递给对象同步模块
            if self.plugin.object_sync:
                for obj_id, obj_data in data.items():
                    self.plugin.object_sync.receive_sync_data(obj_id, obj_data)
                    
            return True
        except Exception as e:
            print(f"✗ 同步数据消息处理失败: {e}")
            return False
            
    def _handle_sync_request(self, data: Dict[str, Any]) -> bool:
        """处理同步请求消息"""
        try:
            print(f"✓ 处理同步请求消息: {data}")
            return True
        except Exception as e:
            print(f"✗ 同步请求消息处理失败: {e}")
            return False
            
    def _handle_sync_response(self, data: Dict[str, Any]) -> bool:
        """处理同步响应消息"""
        try:
            print(f"✓ 处理同步响应消息: {data}")
            return True
        except Exception as e:
            print(f"✗ 同步响应消息处理失败: {e}")
            return False
            
    def _handle_rpc_call(self, data: Dict[str, Any]) -> bool:
        """处理RPC调用消息"""
        try:
            print(f"✓ 处理RPC调用消息: {data}")
            return True
        except Exception as e:
            print(f"✗ RPC调用消息处理失败: {e}")
            return False
            
    def _handle_rpc_response(self, data: Dict[str, Any]) -> bool:
        """处理RPC响应消息"""
        try:
            print(f"✓ 处理RPC响应消息: {data}")
            return True
        except Exception as e:
            print(f"✗ RPC响应消息处理失败: {e}")
            return False
            
    def _handle_event(self, data: Dict[str, Any]) -> bool:
        """处理事件消息"""
        try:
            print(f"✓ 处理事件消息: {data}")
            return True
        except Exception as e:
            print(f"✗ 事件消息处理失败: {e}")
            return False
            
    def _handle_chat(self, data: Dict[str, Any]) -> bool:
        """处理聊天消息"""
        try:
            print(f"✓ 处理聊天消息: {data}")
            return True
        except Exception as e:
            print(f"✗ 聊天消息处理失败: {e}")
            return False
            
    def _handle_error(self, data: Dict[str, Any]) -> bool:
        """处理错误消息"""
        try:
            print(f"✗ 处理错误消息: {data}")
            
            # 触发协议错误回调
            self._trigger_protocol_callback('protocol_error', data)
            
            return True
        except Exception as e:
            print(f"✗ 错误消息处理失败: {e}")
            return False
            
    def get_protocol_config(self) -> Dict[str, Any]:
        """
        获取协议配置
        
        Returns:
            协议配置字典
        """
        return self.protocol_config.copy()
        
    def set_protocol_config(self, config: Dict[str, Any]) -> bool:
        """
        设置协议配置
        
        Args:
            config: 协议配置字典
            
        Returns:
            是否设置成功
        """
        try:
            self.protocol_config.update(config)
            print(f"✓ 协议配置已更新: {self.protocol_config}")
            return True
        except Exception as e:
            print(f"✗ 协议配置设置失败: {e}")
            return False
            
    def get_protocol_stats(self) -> Dict[str, Any]:
        """
        获取协议统计信息
        
        Returns:
            协议统计字典
        """
        return self.protocol_stats.copy()
        
    def reset_protocol_stats(self):
        """重置协议统计信息"""
        try:
            self.protocol_stats = {
                'messages_sent': 0,
                'messages_received': 0,
                'bytes_sent': 0,
                'bytes_received': 0,
                'compressed_messages': 0,
                'checksum_errors': 0,
                'protocol_errors': 0,
                'fragmented_packets': 0
            }
            print("✓ 协议统计信息已重置")
        except Exception as e:
            print(f"✗ 协议统计信息重置失败: {e}")
            
    def get_protocol_state(self) -> Dict[str, Any]:
        """
        获取协议状态
        
        Returns:
            协议状态字典
        """
        return self.protocol_state.copy()
        
    def _trigger_protocol_callback(self, callback_type: str, data: Dict[str, Any]):
        """
        触发协议回调
        
        Args:
            callback_type: 回调类型
            data: 回调数据
        """
        try:
            if callback_type in self.protocol_callbacks:
                for callback in self.protocol_callbacks[callback_type]:
                    try:
                        callback(data)
                    except Exception as e:
                        print(f"✗ 协议回调执行失败: {e}")
        except Exception as e:
            print(f"✗ 协议回调触发失败: {e}")
            
    def register_protocol_callback(self, callback_type: str, callback: callable):
        """
        注册协议回调
        
        Args:
            callback_type: 回调类型
            callback: 回调函数
        """
        try:
            if callback_type in self.protocol_callbacks:
                self.protocol_callbacks[callback_type].append(callback)
                print(f"✓ 协议回调已注册: {callback_type}")
            else:
                print(f"✗ 无效的回调类型: {callback_type}")
        except Exception as e:
            print(f"✗ 协议回调注册失败: {e}")
            
    def unregister_protocol_callback(self, callback_type: str, callback: callable):
        """
        注销协议回调
        
        Args:
            callback_type: 回调类型
            callback: 回调函数
        """
        try:
            if callback_type in self.protocol_callbacks:
                if callback in self.protocol_callbacks[callback_type]:
                    self.protocol_callbacks[callback_type].remove(callback)
                    print(f"✓ 协议回调已注销: {callback_type}")
        except Exception as e:
            print(f"✗ 协议回调注销失败: {e}")