# -*- coding: utf-8 -*-
"""
AIfeng/2025-01-27 16:30:00
WSA WebSocket服务
将原有wsa_server功能集成到统一WebSocket架构中
"""

import asyncio
import json
import queue
from typing import Dict, Any, Optional, Set
from threading import Lock
from aiohttp import web

from .websocket_service_base import WebSocketServiceBase
from .unified_websocket_manager import WebSocketSession


class WSAWebSocketService(WebSocketServiceBase):
    """WSA WebSocket服务
    
    提供与原wsa_server兼容的功能：
    - Web连接管理
    - Human连接管理  
    - 命令队列处理
    - 消息转发
    """
    
    def __init__(self, manager):
        super().__init__("wsa")
        
        # 连接管理
        self._web_connections: Dict[str, Set[WebSocketSession]] = {}
        self._human_connections: Dict[str, Set[WebSocketSession]] = {}
        self._connection_lock = Lock()
        
        # 命令队列
        self._web_command_queue = queue.Queue()
        self._human_command_queue = queue.Queue()
        
        # 后台任务
        self._queue_processor_task: Optional[asyncio.Task] = None
        
    async def _register_message_handlers(self):
        """注册消息处理器"""
        self.manager.register_message_handler("wsa_register_web", self._handle_register_web)
        self.manager.register_message_handler("wsa_register_human", self._handle_register_human)
        self.manager.register_message_handler("wsa_unregister", self._handle_unregister)
        self.manager.register_message_handler("wsa_get_status", self._handle_get_status)
        
    async def _start_background_tasks(self):
        """启动后台任务"""
        self._queue_processor_task = asyncio.create_task(self._process_command_queues())
        
    async def _cleanup(self):
        """清理资源"""
        if self._queue_processor_task:
            self._queue_processor_task.cancel()
            try:
                await self._queue_processor_task
            except asyncio.CancelledError:
                pass
                
    async def _on_session_disconnected(self, session: WebSocketSession):
        """会话断开处理"""
        with self._connection_lock:
            # 从web连接中移除
            for username, sessions in list(self._web_connections.items()):
                if session in sessions:
                    sessions.discard(session)
                    if not sessions:
                        del self._web_connections[username]
                        
            # 从human连接中移除
            for username, sessions in list(self._human_connections.items()):
                if session in sessions:
                    sessions.discard(session)
                    if not sessions:
                        del self._human_connections[username]
                        
    async def _handle_register_web(self, websocket: web.WebSocketResponse, data: Dict[str, Any]):
        """注册Web连接"""
        username = data.get('username')
        if not username:
            await websocket.send_str(json.dumps({
                "type": "wsa_error",
                "message": "用户名不能为空"
            }))
            return
            
        session = self.manager.get_session(websocket)
        if not session:
            await websocket.send_str(json.dumps({
                "type": "wsa_error",
                "message": "会话未找到"
            }))
            return
            
        with self._connection_lock:
            if username not in self._web_connections:
                self._web_connections[username] = set()
            self._web_connections[username].add(session)
            
        await websocket.send_str(json.dumps({
            "type": "wsa_registered",
            "connection_type": "web",
            "username": username
        }))
        
    async def _handle_register_human(self, websocket: web.WebSocketResponse, data: Dict[str, Any]):
        """注册Human连接"""
        username = data.get('username')
        if not username:
            await websocket.send_str(json.dumps({
                "type": "wsa_error",
                "message": "用户名不能为空"
            }))
            return
            
        session = self.manager.get_session(websocket)
        if not session:
            await websocket.send_str(json.dumps({
                "type": "wsa_error",
                "message": "会话未找到"
            }))
            return
            
        with self._connection_lock:
            if username not in self._human_connections:
                self._human_connections[username] = set()
            self._human_connections[username].add(session)
            
        await websocket.send_str(json.dumps({
            "type": "wsa_registered",
            "connection_type": "human",
            "username": username
        }))
        
    async def _handle_unregister(self, websocket: web.WebSocketResponse, data: Dict[str, Any]):
        """注销连接"""
        username = data.get('username')
        connection_type = data.get('connection_type', 'both')
        
        session = self.manager.get_session(websocket)
        if not session:
            return
            
        with self._connection_lock:
            if connection_type in ['web', 'both'] and username in self._web_connections:
                self._web_connections[username].discard(session)
                if not self._web_connections[username]:
                    del self._web_connections[username]
                    
            if connection_type in ['human', 'both'] and username in self._human_connections:
                self._human_connections[username].discard(session)
                if not self._human_connections[username]:
                    del self._human_connections[username]
                    
        await websocket.send_str(json.dumps({
            "type": "wsa_unregistered",
            "username": username,
            "connection_type": connection_type
        }))
        
    async def _handle_get_status(self, websocket: web.WebSocketResponse, data: Dict[str, Any]):
        """获取连接状态"""
        with self._connection_lock:
            web_users = list(self._web_connections.keys())
            human_users = list(self._human_connections.keys())
            
        await websocket.send_str(json.dumps({
            "type": "wsa_status",
            "data": {
                "web_connections": len(self._web_connections),
                "human_connections": len(self._human_connections),
                "web_users": web_users,
                "human_users": human_users,
                "web_queue_size": self._web_command_queue.qsize(),
                "human_queue_size": self._human_command_queue.qsize()
            }
        }))
        
    async def _process_command_queues(self):
        """处理命令队列"""
        while True:
            try:
                # 处理Web命令队列
                await self._process_web_commands()
                
                # 处理Human命令队列
                await self._process_human_commands()
                
                # 短暂休眠避免CPU占用过高
                await asyncio.sleep(0.01)
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"命令队列处理错误: {e}")
                await asyncio.sleep(0.1)
                
    async def _process_web_commands(self):
        """处理Web命令队列"""
        try:
            while True:
                try:
                    command = self._web_command_queue.get_nowait()
                    await self._forward_web_command(command)
                except queue.Empty:
                    break
        except Exception as e:
            self.logger.error(f"Web命令处理错误: {e}")
            
    async def _process_human_commands(self):
        """处理Human命令队列"""
        try:
            while True:
                try:
                    command = self._human_command_queue.get_nowait()
                    await self._forward_human_command(command)
                except queue.Empty:
                    break
        except Exception as e:
            self.logger.error(f"Human命令处理错误: {e}")
            
    async def _forward_web_command(self, command: Dict[str, Any]):
        """转发Web命令"""
        username = command.get('Username')
        if not username:
            return
            
        with self._connection_lock:
            sessions = self._web_connections.get(username, set())
            
        if sessions:
            message = {
                "type": "wsa_command",
                "source": "web",
                "data": command
            }
            
            for session in list(sessions):
                try:
                    await session.send_message(message)
                except Exception as e:
                    self.logger.error(f"发送Web命令失败 [{username}]: {e}")
                    
    async def _forward_human_command(self, command: Dict[str, Any]):
        """转发Human命令"""
        username = command.get('Username')
        if not username:
            return
            
        with self._connection_lock:
            sessions = self._human_connections.get(username, set())
            
        if sessions:
            message = {
                "type": "wsa_command",
                "source": "human",
                "data": command
            }
            
            for session in list(sessions):
                try:
                    await session.send_message(message)
                except Exception as e:
                    self.logger.error(f"发送Human命令失败 [{username}]: {e}")
                    
    # 兼容性接口
    def is_connected(self, username: str) -> bool:
        """检查Web用户是否已连接"""
        with self._connection_lock:
            return username in self._web_connections and bool(self._web_connections[username])
            
    def is_connected_human(self, username: str) -> bool:
        """检查Human用户是否已连接"""
        with self._connection_lock:
            return username in self._human_connections and bool(self._human_connections[username])
            
    def add_connection(self, username: str, connection: Any):
        """添加连接（兼容性接口，已废弃）"""
        self.logger.warning("add_connection方法已废弃，请使用消息注册机制")
        
    def remove_connection(self, username: str):
        """移除连接（兼容性接口，已废弃）"""
        self.logger.warning("remove_connection方法已废弃，连接会自动清理")
        
    def add_cmd(self, command: Dict[str, Any], target: str = "web"):
        """添加命令到队列"""
        try:
            if target == "web":
                self._web_command_queue.put(command, timeout=1.0)
            elif target == "human":
                self._human_command_queue.put(command, timeout=1.0)
            else:
                self.logger.warning(f"未知的目标类型: {target}")
        except queue.Full:
            self.logger.warning(f"命令队列已满，丢弃命令: {command}")
            
    async def send_direct_message(self, message: Dict[str, Any], target: str = "web"):
        """直接发送消息（不封装为wsa_command）"""
        username = message.get('Username')
        if not username:
            self.logger.warning("消息缺少Username字段")
            return
            
        with self._connection_lock:
            if target == "web":
                sessions = self._web_connections.get(username, set())
            elif target == "human":
                sessions = self._human_connections.get(username, set())
            else:
                self.logger.warning(f"未知的目标类型: {target}")
                return
                
        if sessions:
            for session in list(sessions):
                try:
                    await session.send_message(message)
                except Exception as e:
                    self.logger.error(f"直接发送消息失败 [{username}]: {e}")
        else:
            self.logger.debug(f"用户 {username} 未连接，无法发送直接消息")
            
    def get_cmd(self, timeout: float = 1.0, target: str = "web") -> Optional[Dict[str, Any]]:
        """从队列获取命令"""
        try:
            if target == "web":
                return self._web_command_queue.get(timeout=timeout)
            elif target == "human":
                return self._human_command_queue.get(timeout=timeout)
            else:
                self.logger.warning(f"未知的目标类型: {target}")
                return None
        except queue.Empty:
            return None
            
    def get_connection_count(self, target: str = "web") -> int:
        """获取连接数量"""
        with self._connection_lock:
            if target == "web":
                return len(self._web_connections)
            elif target == "human":
                return len(self._human_connections)
            else:
                return len(self._web_connections) + len(self._human_connections)
                
    def get_usernames(self, target: str = "web") -> list:
        """获取用户名列表"""
        with self._connection_lock:
            if target == "web":
                return list(self._web_connections.keys())
            elif target == "human":
                return list(self._human_connections.keys())
            else:
                return list(set(list(self._web_connections.keys()) + list(self._human_connections.keys())))


# 兼容性包装器
class WSAWebSocketManager:
    """WSA WebSocket管理器兼容性包装器"""
    
    def __init__(self, service: WSAWebSocketService):
        self.service = service
        
    def is_connected(self, username: str) -> bool:
        return self.service.is_connected(username)
        
    def is_connected_human(self, username: str) -> bool:
        return self.service.is_connected_human(username)
        
    def add_connection(self, username: str, connection: Any):
        self.service.add_connection(username, connection)
        
    def remove_connection(self, username: str):
        self.service.remove_connection(username)
        
    def add_cmd(self, command: Dict[str, Any]):
        self.service.add_cmd(command, "web")
        
    async def send_direct_message(self, message: Dict[str, Any]):
        """直接发送消息（不封装为wsa_command）"""
        await self.service.send_direct_message(message, "web")
        
    def get_cmd(self, timeout: float = 1.0) -> Optional[Dict[str, Any]]:
        return self.service.get_cmd(timeout, "web")
        
    def get_connection_count(self) -> int:
        return self.service.get_connection_count("web")
        
    def get_usernames(self) -> list:
        return self.service.get_usernames("web")


# 全局实例（兼容性）
_wsa_service: Optional[WSAWebSocketService] = None
_web_instance: Optional[WSAWebSocketManager] = None
_human_instance: Optional[WSAWebSocketManager] = None


def initialize_wsa_service(service: WSAWebSocketService):
    """初始化WSA服务"""
    global _wsa_service, _web_instance, _human_instance
    _wsa_service = service
    _web_instance = WSAWebSocketManager(service)
    _human_instance = WSAWebSocketManager(service)


def get_web_instance() -> WSAWebSocketManager:
    """获取Web WebSocket管理器实例"""
    if _web_instance is None:
        raise RuntimeError("WSA服务未初始化")
    return _web_instance


def get_instance() -> WSAWebSocketManager:
    """获取Human WebSocket管理器实例"""
    if _human_instance is None:
        raise RuntimeError("WSA服务未初始化")
    return _human_instance