"""
WebSocket服务器模块
处理实时双向通信连接
"""

import time
import threading
import json
import base64
from typing import Dict, Any, List, Optional

class WebSocketServer:
    """
    WebSocket服务器
    处理实时双向通信连接
    """
    
    def __init__(self, plugin):
        """
        初始化WebSocket服务器
        
        Args:
            plugin: 实时通信插件实例
        """
        self.plugin = plugin
        self.enabled = False
        self.initialized = False
        
        # WebSocket配置
        self.websocket_config = {
            "host": "0.0.0.0",
            "port": 8080,
            "enable_ssl": False,
            "ssl_cert_file": "",
            "ssl_key_file": "",
            "max_connections": 1000,
            "enable_compression": True,
            "ping_interval": 30.0,
            "ping_timeout": 10.0,
            "message_queue_size": 1000,
            "enable_cors": True
        }
        
        # WebSocket服务器状态
        self.server_state = {
            "is_running": False,
            "start_time": 0.0,
            "uptime": 0.0,
            "current_connections": 0,
            "total_connections": 0,
            "bytes_sent": 0,
            "bytes_received": 0
        }
        
        # WebSocket服务器统计
        self.server_stats = {
            "messages_sent": 0,
            "messages_received": 0,
            "connections_accepted": 0,
            "connections_dropped": 0,
            "errors": 0
        }
        
        # 连接存储
        self.connections = {}
        self.connection_lock = threading.RLock()
        
        # 消息队列
        self.message_queue = []
        self.message_queue_lock = threading.RLock()
        
        # 服务器线程
        self.server_thread = None
        self.server_thread_running = False
        
        # 回调函数
        self.websocket_callbacks = {
            "server_started": [],
            "server_stopped": [],
            "client_connected": [],
            "client_disconnected": [],
            "message_received": [],
            "websocket_error": []
        }
        
        # 时间戳记录
        self.last_ping_time = 0.0
        self.last_message_time = 0.0
        
        print("✓ WebSocket服务器已创建")
        
    def initialize(self) -> bool:
        """
        初始化WebSocket服务器
        
        Returns:
            是否初始化成功
        """
        try:
            print("正在初始化WebSocket服务器...")
            
            self.initialized = True
            print("✓ WebSocket服务器初始化完成")
            return True
            
        except Exception as e:
            print(f"✗ WebSocket服务器初始化失败: {e}")
            self.server_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            return False
            
    def enable(self) -> bool:
        """
        启用WebSocket服务器
        
        Returns:
            是否启用成功
        """
        try:
            if not self.initialized:
                print("✗ WebSocket服务器未初始化")
                return False
                
            # 启动WebSocket服务器
            self._start_websocket_server()
            
            self.enabled = True
            print("✓ WebSocket服务器已启用")
            return True
            
        except Exception as e:
            print(f"✗ WebSocket服务器启用失败: {e}")
            self.server_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            return False
            
    def disable(self):
        """禁用WebSocket服务器"""
        try:
            self.enabled = False
            
            # 停止WebSocket服务器
            self._stop_websocket_server()
            
            print("✓ WebSocket服务器已禁用")
            
        except Exception as e:
            print(f"✗ WebSocket服务器禁用失败: {e}")
            self.server_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            
    def finalize(self):
        """清理WebSocket服务器资源"""
        try:
            # 禁用WebSocket服务器
            if self.enabled:
                self.disable()
                
            # 清理回调
            self.websocket_callbacks.clear()
            
            self.initialized = False
            print("✓ WebSocket服务器资源已清理")
            
        except Exception as e:
            print(f"✗ WebSocket服务器资源清理失败: {e}")
            import traceback
            traceback.print_exc()
            
    def update(self, dt: float):
        """
        更新WebSocket服务器状态
        
        Args:
            dt: 时间增量（秒）
        """
        try:
            if not self.enabled:
                return
                
            current_time = time.time()
            
            # 更新运行时间
            if self.server_state["is_running"]:
                self.server_state["uptime"] = current_time - self.server_state["start_time"]
                
            # 处理消息队列
            self._process_message_queue()
            
            # 发送心跳
            if current_time - self.last_ping_time > self.websocket_config["ping_interval"]:
                self._send_ping_to_all_clients()
                self.last_ping_time = current_time
                
            self.last_message_time = current_time
            
        except Exception as e:
            print(f"✗ WebSocket服务器更新失败: {e}")
            self.server_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            
    def _start_websocket_server(self):
        """启动WebSocket服务器"""
        try:
            if self.server_state["is_running"]:
                return
                
            # 在实际实现中，这里会启动一个WebSocket服务器
            # 由于这是一个示例，我们模拟服务器启动过程
            self.server_state["is_running"] = True
            self.server_state["start_time"] = time.time()
            
            # 触发服务器启动回调
            self._trigger_websocket_callback("server_started", {
                "host": self.websocket_config["host"],
                "port": self.websocket_config["port"],
                "timestamp": self.server_state["start_time"]
            })
            
            print(f"✓ WebSocket服务器已启动: ws://{self.websocket_config['host']}:{self.websocket_config['port']}")
            
        except Exception as e:
            print(f"✗ WebSocket服务器启动失败: {e}")
            self.server_stats["errors"] += 1
            
    def _stop_websocket_server(self):
        """停止WebSocket服务器"""
        try:
            if not self.server_state["is_running"]:
                return
                
            # 断开所有客户端连接
            self._disconnect_all_clients()
            
            self.server_state["is_running"] = False
            self.server_state["start_time"] = 0.0
            self.server_state["uptime"] = 0.0
            
            # 触发服务器停止回调
            self._trigger_websocket_callback("server_stopped", {
                "timestamp": time.time()
            })
            
            print("✓ WebSocket服务器已停止")
            
        except Exception as e:
            print(f"✗ WebSocket服务器停止失败: {e}")
            self.server_stats["errors"] += 1
            
    def _process_message_queue(self):
        """处理消息队列"""
        try:
            if not self.message_queue:
                return
                
            with self.message_queue_lock:
                messages_to_process = self.message_queue[:]
                self.message_queue.clear()
                
            for message_data in messages_to_process:
                try:
                    client_id = message_data.get("client_id")
                    message = message_data.get("message")
                    
                    # 处理消息
                    self._handle_incoming_message(client_id, message)
                    
                except Exception as e:
                    print(f"✗ 消息处理失败: {e}")
                    self.server_stats["errors"] += 1
                    
        except Exception as e:
            print(f"✗ 消息队列处理失败: {e}")
            self.server_stats["errors"] += 1
            
    def _handle_incoming_message(self, client_id: str, message: Dict[str, Any]):
        """
        处理传入消息
        
        Args:
            client_id: 客户端ID
            message: 消息数据
        """
        try:
            self.server_stats["messages_received"] += 1
            
            # 触发消息接收回调
            self._trigger_websocket_callback("message_received", {
                "client_id": client_id,
                "message": message,
                "timestamp": time.time()
            })
            
            # 将消息传递给消息路由器处理
            if self.plugin.message_router:
                self.plugin.message_router.route_message(client_id, message)
                
        except Exception as e:
            print(f"✗ 传入消息处理失败: {e}")
            self.server_stats["errors"] += 1
            
    def _send_ping_to_all_clients(self):
        """向所有客户端发送心跳"""
        try:
            ping_message = {
                "type": "ping",
                "timestamp": time.time()
            }
            
            with self.connection_lock:
                for client_id, connection in self.connections.items():
                    try:
                        self._send_message_to_client(client_id, ping_message)
                    except Exception as e:
                        print(f"✗ 心跳发送失败: {client_id} - {e}")
                        
        except Exception as e:
            print(f"✗ 心跳发送失败: {e}")
            self.server_stats["errors"] += 1
            
    def _disconnect_all_clients(self):
        """断开所有客户端连接"""
        try:
            with self.connection_lock:
                client_ids = list(self.connections.keys())
                
            for client_id in client_ids:
                self.disconnect_client(client_id, "server_shutdown")
                
        except Exception as e:
            print(f"✗ 断开所有客户端连接失败: {e}")
            self.server_stats["errors"] += 1
            
    def send_message_to_client(self, client_id: str, message: Dict[str, Any]) -> bool:
        """
        向客户端发送消息
        
        Args:
            client_id: 客户端ID
            message: 消息数据
            
        Returns:
            是否发送成功
        """
        try:
            if not self.enabled:
                return False
                
            return self._send_message_to_client(client_id, message)
            
        except Exception as e:
            print(f"✗ 向客户端发送消息失败: {e}")
            self.server_stats["errors"] += 1
            return False
            
    def _send_message_to_client(self, client_id: str, message: Dict[str, Any]) -> bool:
        """
        向客户端发送消息（内部实现）
        
        Args:
            client_id: 客户端ID
            message: 消息数据
            
        Returns:
            是否发送成功
        """
        try:
            # 检查客户端连接是否存在
            with self.connection_lock:
                if client_id not in self.connections:
                    return False
                    
            # 序列化消息
            if self.plugin.data_serializer:
                serialized_message = self.plugin.data_serializer.serialize(message)
            else:
                serialized_message = json.dumps(message, ensure_ascii=False)
                
            # 发送消息（模拟）
            message_size = len(serialized_message.encode('utf-8'))
            self.server_state["bytes_sent"] += message_size
            self.server_stats["messages_sent"] += 1
            
            return True
            
        except Exception as e:
            print(f"✗ 向客户端发送消息失败: {e}")
            self.server_stats["errors"] += 1
            return False
            
    def broadcast_message(self, message: Dict[str, Any], exclude_client_id: str = None) -> bool:
        """
        广播消息到所有客户端
        
        Args:
            message: 消息数据
            exclude_client_id: 要排除的客户端ID
            
        Returns:
            是否广播成功
        """
        try:
            if not self.enabled:
                return False
                
            success_count = 0
            with self.connection_lock:
                for client_id in self.connections.keys():
                    # 跳过排除的客户端
                    if client_id == exclude_client_id:
                        continue
                        
                    if self._send_message_to_client(client_id, message):
                        success_count += 1
                        
            return success_count > 0
            
        except Exception as e:
            print(f"✗ 消息广播失败: {e}")
            self.server_stats["errors"] += 1
            return False
            
    def add_client_connection(self, client_id: str, connection_data: Dict[str, Any]) -> bool:
        """
        添加客户端连接
        
        Args:
            client_id: 客户端ID
            connection_data: 连接数据
            
        Returns:
            是否添加成功
        """
        try:
            if not self.enabled:
                return False
                
            with self.connection_lock:
                # 检查连接数限制
                if len(self.connections) >= self.websocket_config["max_connections"]:
                    print("✗ 连接数已达上限")
                    return False
                    
                # 添加连接
                self.connections[client_id] = {
                    "connection_data": connection_data,
                    "connect_time": time.time(),
                    "last_activity": time.time(),
                    "bytes_sent": 0,
                    "bytes_received": 0,
                    "messages_sent": 0,
                    "messages_received": 0
                }
                
            self.server_state["current_connections"] = len(self.connections)
            self.server_state["total_connections"] += 1
            self.server_stats["connections_accepted"] += 1
            
            # 触发客户端连接回调
            self._trigger_websocket_callback("client_connected", {
                "client_id": client_id,
                "connection_data": connection_data,
                "timestamp": time.time()
            })
            
            print(f"✓ 客户端已连接: {client_id}")
            return True
            
        except Exception as e:
            print(f"✗ 添加客户端连接失败: {e}")
            self.server_stats["errors"] += 1
            return False
            
    def remove_client_connection(self, client_id: str, reason: str = "unknown") -> bool:
        """
        移除客户端连接
        
        Args:
            client_id: 客户端ID
            reason: 断开原因
            
        Returns:
            是否移除成功
        """
        try:
            with self.connection_lock:
                if client_id in self.connections:
                    del self.connections[client_id]
                    
            self.server_state["current_connections"] = len(self.connections)
            self.server_stats["connections_dropped"] += 1
            
            # 触发客户端断开回调
            self._trigger_websocket_callback("client_disconnected", {
                "client_id": client_id,
                "reason": reason,
                "timestamp": time.time()
            })
            
            print(f"✓ 客户端已断开: {client_id} (原因: {reason})")
            return True
            
        except Exception as e:
            print(f"✗ 移除客户端连接失败: {e}")
            self.server_stats["errors"] += 1
            return False
            
    def disconnect_client(self, client_id: str, reason: str = "kicked") -> bool:
        """
        断开客户端连接
        
        Args:
            client_id: 客户端ID
            reason: 断开原因
            
        Returns:
            是否断开成功
        """
        try:
            # 通知客户端断开连接
            disconnect_message = {
                "type": "disconnect",
                "reason": reason,
                "timestamp": time.time()
            }
            
            self._send_message_to_client(client_id, disconnect_message)
            
            # 移除连接
            return self.remove_client_connection(client_id, reason)
            
        except Exception as e:
            print(f"✗ 断开客户端连接失败: {e}")
            self.server_stats["errors"] += 1
            return False
            
    def get_connection_info(self, client_id: str) -> Optional[Dict[str, Any]]:
        """
        获取连接信息
        
        Args:
            client_id: 客户端ID
            
        Returns:
            连接信息或None
        """
        try:
            with self.connection_lock:
                if client_id in self.connections:
                    return self.connections[client_id].copy()
                else:
                    return None
                    
        except Exception as e:
            print(f"✗ 获取连接信息失败: {e}")
            self.server_stats["errors"] += 1
            return None
            
    def get_all_connections(self) -> Dict[str, Any]:
        """
        获取所有连接信息
        
        Returns:
            所有连接信息字典
        """
        try:
            with self.connection_lock:
                return {k: v.copy() for k, v in self.connections.items()}
        except Exception as e:
            print(f"✗ 获取所有连接信息失败: {e}")
            self.server_stats["errors"] += 1
            return {}
            
    def get_server_stats(self) -> Dict[str, Any]:
        """
        获取服务器统计信息
        
        Returns:
            服务器统计字典
        """
        return {
            "state": self.server_state.copy(),
            "stats": self.server_stats.copy(),
            "config": self.websocket_config.copy()
        }
        
    def reset_stats(self):
        """重置服务器统计信息"""
        try:
            self.server_stats = {
                "messages_sent": 0,
                "messages_received": 0,
                "connections_accepted": 0,
                "connections_dropped": 0,
                "errors": 0
            }
            
            self.server_state["bytes_sent"] = 0
            self.server_state["bytes_received"] = 0
            
            print("✓ 服务器统计信息已重置")
        except Exception as e:
            print(f"✗ 服务器统计信息重置失败: {e}")
            
    def set_websocket_config(self, config: Dict[str, Any]) -> bool:
        """
        设置WebSocket配置
        
        Args:
            config: WebSocket配置字典
            
        Returns:
            是否设置成功
        """
        try:
            old_port = self.websocket_config.get("port")
            new_port = config.get("port", old_port)
            
            # 如果服务器正在运行且端口发生变化，重启服务器
            if (self.server_state["is_running"] and 
                old_port != new_port):
                self._stop_websocket_server()
                self.websocket_config.update(config)
                self._start_websocket_server()
            else:
                self.websocket_config.update(config)
                
            print(f"✓ WebSocket配置已更新: {self.websocket_config}")
            return True
        except Exception as e:
            print(f"✗ WebSocket配置设置失败: {e}")
            return False
            
    def get_websocket_config(self) -> Dict[str, Any]:
        """
        获取WebSocket配置
        
        Returns:
            WebSocket配置字典
        """
        return self.websocket_config.copy()
        
    def _trigger_websocket_callback(self, callback_type: str, data: Dict[str, Any]):
        """
        触发WebSocket回调
        
        Args:
            callback_type: 回调类型
            data: 回调数据
        """
        try:
            if callback_type in self.websocket_callbacks:
                for callback in self.websocket_callbacks[callback_type]:
                    try:
                        callback(data)
                    except Exception as e:
                        print(f"✗ WebSocket回调执行失败: {callback_type} - {e}")
        except Exception as e:
            print(f"✗ WebSocket回调触发失败: {e}")
            
    def register_websocket_callback(self, callback_type: str, callback: callable):
        """
        注册WebSocket回调
        
        Args:
            callback_type: 回调类型
            callback: 回调函数
        """
        try:
            if callback_type in self.websocket_callbacks:
                self.websocket_callbacks[callback_type].append(callback)
                print(f"✓ WebSocket回调已注册: {callback_type}")
            else:
                print(f"✗ 无效的回调类型: {callback_type}")
        except Exception as e:
            print(f"✗ WebSocket回调注册失败: {e}")
            
    def unregister_websocket_callback(self, callback_type: str, callback: callable):
        """
        注销WebSocket回调
        
        Args:
            callback_type: 回调类型
            callback: 回调函数
        """
        try:
            if callback_type in self.websocket_callbacks:
                if callback in self.websocket_callbacks[callback_type]:
                    self.websocket_callbacks[callback_type].remove(callback)
                    print(f"✓ WebSocket回调已注销: {callback_type}")
            else:
                print(f"✗ 无效的回调类型: {callback_type}")
        except Exception as e:
            print(f"✗ WebSocket回调注销失败: {e}")