# -*- coding: utf-8 -*-
"""
AIfeng/2025-07-15 14:41:21
统一WebSocket管理模块
提供统一的WebSocket连接管理、消息推送和事件处理功能
"""

import json
import time
import asyncio
import weakref
from typing import Dict, Any, Optional, Callable, List, Set
from threading import Lock
from aiohttp import web, WSMsgType
from logger import logger


class WebSocketSession:
    """WebSocket会话管理类"""
    
    def __init__(self, session_id: str, websocket: web.WebSocketResponse):
        self.session_id = session_id
        self.websocket = websocket
        self.created_at = time.time()
        self.last_ping = time.time()
        self.metadata = {}
        
    def __eq__(self, other):
        """基于session_id判断会话是否相等"""
        if not isinstance(other, WebSocketSession):
            return False
        return self.session_id == other.session_id
        
    def __hash__(self):
        """基于session_id生成哈希值"""
        return hash(self.session_id)
        
    def is_alive(self) -> bool:
        """检查连接是否存活"""
        return not self.websocket.closed
        
    async def send_message(self, message: Dict[str, Any]) -> bool:
        """发送消息到WebSocket客户端"""
        try:
            if not self.is_alive():
                return False
            await self.websocket.send_str(json.dumps(message))
            return True
        except ConnectionResetError:
            logger.warning(f'[Session:{self.session_id}] 客户端连接已重置')
            return False
        except ConnectionAbortedError:
            logger.warning(f'[Session:{self.session_id}] 客户端连接已中止')
            return False
        except Exception as e:
            logger.error(f'[Session:{self.session_id}] 发送消息失败: {e}')
            return False
            
    def update_ping(self):
        """更新心跳时间"""
        self.last_ping = time.time()
        
    def set_metadata(self, key: str, value: Any):
        """设置会话元数据"""
        self.metadata[key] = value
        
    def get_metadata(self, key: str, default=None):
        """获取会话元数据"""
        return self.metadata.get(key, default)
        
    async def close(self):
        """关闭WebSocket连接"""
        try:
            if not self.websocket.closed:
                await self.websocket.close()
                logger.info(f'[Session:{self.session_id}] WebSocket连接已关闭')
        except ConnectionResetError:
            logger.warning(f'[Session:{self.session_id}] 连接已被远程主机重置，无需关闭')
        except ConnectionAbortedError:
            logger.warning(f'[Session:{self.session_id}] 连接已被中止，无需关闭')
        except Exception as e:
            logger.error(f'[Session:{self.session_id}] 关闭WebSocket连接失败: {e}')


class UnifiedWebSocketManager:
    """统一WebSocket管理器"""
    
    def __init__(self):
        self._sessions: Dict[str, WebSocketSession] = {}  # session_id -> WebSocketSession (单连接)
        self._websockets: Dict[web.WebSocketResponse, WebSocketSession] = {}  # websocket -> session映射
        self._message_handlers: Dict[str, Callable] = {}  # 消息类型处理器
        self._event_handlers: Dict[str, List[Callable]] = {}  # 事件处理器
        self._lock = Lock()
        
        # 注册默认消息处理器
        self.register_message_handler('ping', self._handle_ping)
        self.register_message_handler('login', self._handle_login)
        
    def register_message_handler(self, message_type: str, handler: Callable):
        """注册消息处理器"""
        self._message_handlers[message_type] = handler
        logger.info(f'注册消息处理器: {message_type}')
        
    def register_event_handler(self, event_type: str, handler: Callable):
        """注册事件处理器"""
        if event_type not in self._event_handlers:
            self._event_handlers[event_type] = []
        self._event_handlers[event_type].append(handler)
        logger.info(f'注册事件处理器: {event_type}')
        
    async def _emit_event(self, event_type: str, **kwargs):
        """触发事件"""
        if event_type in self._event_handlers:
            for handler in self._event_handlers[event_type]:
                try:
                    if asyncio.iscoroutinefunction(handler):
                        await handler(**kwargs)
                    else:
                        handler(**kwargs)
                except Exception as e:
                    logger.error(f'事件处理器执行失败 {event_type}: {e}')
                    
    def add_session(self, session_id: str, websocket: web.WebSocketResponse) -> WebSocketSession:
        """添加WebSocket会话 - 单连接模式，新连接会替换旧连接"""
        with self._lock:
            # 检查是否已存在相同的websocket连接
            if websocket in self._websockets:
                existing_session = self._websockets[websocket]
                logger.warning(f'[Session:{session_id}] WebSocket连接已存在 (WebSocket={id(websocket)}, 原Session={existing_session.session_id})')
                return existing_session
                
            session = WebSocketSession(session_id, websocket)
            
            # 如果已存在该session_id的连接，先清理旧连接
            if session_id in self._sessions:
                old_session = self._sessions[session_id]
                # 从websockets映射中移除旧连接
                if old_session.websocket in self._websockets:
                    del self._websockets[old_session.websocket]
                # 尝试关闭旧连接
                try:
                    if not old_session.websocket.closed:
                        asyncio.create_task(old_session.websocket.close())
                        logger.info(f"[Session:{session_id}] 已关闭旧连接")
                except Exception as e:
                    logger.warning(f"[Session:{session_id}] 关闭旧连接时出错: {e}")
            
            # 添加新连接 (单连接模式)
            self._sessions[session_id] = session
            self._websockets[websocket] = session
            
            logger.info(f'[Session:{session_id}] 添加WebSocket会话 (单连接模式) (WebSocket={id(websocket)})')
                
            return session
            
    def remove_session(self, websocket: web.WebSocketResponse):
        """移除WebSocket会话"""
        with self._lock:
            if websocket in self._websockets:
                session = self._websockets[websocket]
                session_id = session.session_id
                
                # 从会话映射中移除（单连接模式）
                if session_id in self._sessions and self._sessions[session_id] == session:
                    del self._sessions[session_id]
                        
                # 从websocket映射中移除
                del self._websockets[websocket]
                
                logger.info(f'[Session:{session_id}] 移除WebSocket会话')
                return session
        return None
        
    def get_session(self, websocket: web.WebSocketResponse) -> Optional[WebSocketSession]:
        """获取WebSocket会话"""
        return self._websockets.get(websocket)
        
    def get_sessions_by_id(self, session_id: str) -> Optional[WebSocketSession]:
        """根据会话ID获取WebSocket会话（单连接模式）"""
        with self._lock:
            # 尝试使用原始session_id查找
            session = self._sessions.get(session_id)
            if session:
                return session
            
            # 如果是字符串类型但存储的是整数类型，尝试转换
            if isinstance(session_id, str) and session_id.isdigit():
                int_session_id = int(session_id)
                session = self._sessions.get(int_session_id)
                if session:
                    return session
            
            # 如果是整数类型但存储的是字符串类型，尝试转换
            elif isinstance(session_id, int):
                str_session_id = str(session_id)
                session = self._sessions.get(str_session_id)
                if session:
                    return session
            
            return None
            
    def _update_session_id(self, websocket: web.WebSocketResponse, old_session_id: str, new_session_id: str):
        """更新WebSocket会话的session_id（单连接模式）"""
        with self._lock:
            if websocket in self._websockets:
                session = self._websockets[websocket]
                
                # 从旧的session_id映射中移除（只有当前session匹配时）
                if old_session_id in self._sessions and self._sessions[old_session_id] == session:
                    del self._sessions[old_session_id]
                
                # 更新session的session_id
                session.session_id = new_session_id
                
                # 如果新session_id已存在，先清理旧连接
                if new_session_id in self._sessions:
                    old_session = self._sessions[new_session_id]
                    if old_session.websocket in self._websockets:
                        del self._websockets[old_session.websocket]
                    try:
                        if not old_session.websocket.closed:
                            asyncio.create_task(old_session.websocket.close())
                            logger.info(f"[Session:{new_session_id}] 更新时关闭旧连接")
                    except Exception as e:
                        logger.warning(f"[Session:{new_session_id}] 更新时关闭旧连接出错: {e}")
                
                # 添加到新的session_id映射（单连接模式）
                self._sessions[new_session_id] = session
                
                logger.info(f'[Session] 更新会话ID: {old_session_id} -> {new_session_id}')
                return True
        return False

    async def broadcast_raw_message_to_session(self, session_id: str, message: Dict,source: str = "原数据") -> int:
        """直接广播原始消息到指定会话的所有WebSocket连接"""
        # 确保session_id为字符串类型，保持一致性
        if isinstance(session_id, int):
            session_id = str(session_id)
        elif not isinstance(session_id, str):
            session_id = str(session_id)
            
        session = self.get_sessions_by_id(session_id)
        if not session:
            logger.warning(f'[Session:{session_id}] 没有找到WebSocket连接')
            return 0
            
        # 详细调试日志：显示会话详情
        logger.info(f'[Session:{session_id}] 开始广播消息，找到 1 个连接')
        logger.info(f'[Session:{session_id}] 连接: WebSocket={id(session.websocket)}, 创建时间={session.created_at}, 存活状态={session.is_alive()}')
            
        success_count = 0
        
        logger.info(f'[Session:{session_id}] 正在发送消息 (WebSocket={id(session.websocket)})')
        if await session.send_message(message):
            success_count = 1
            logger.info(f'[Session:{session_id}] 发送成功')
        else:
            logger.warning(f'[Session:{session_id}] 发送失败')
            # 清理失败的连接
            self.remove_session(session.websocket)
            
        logger.info(f'[Session:{session_id}] 广播原始消息完成: 成功{success_count}/总计1')
        return success_count
            
    async def broadcast_to_session(self, session_id: str, message_type: str, content: Any, 
                                 source: str = "系统", metadata: Dict = None) -> int:
        """向指定会话的所有WebSocket连接广播消息"""
        # 确保session_id为字符串类型，保持一致性
        if isinstance(session_id, int):
            session_id = str(session_id)
        elif not isinstance(session_id, str):
            session_id = str(session_id)
            
        session = self.get_sessions_by_id(session_id)
        if not session:
            logger.warning(f'[Session:{session_id}] 没有找到WebSocket连接')
            return 0

        # 详细调试日志：显示会话详情
        logger.info(f'[Session:{session_id}] 开始广播消息，找到 1 个连接')
        logger.info(f'[Session:{session_id}] 连接: WebSocket={id(session.websocket)}, 创建时间={session.created_at}, 存活状态={session.is_alive()}')
            

        message = {
            "type": message_type,
            "session_id": session_id,
            "content": content,
            "source": source,
            "timestamp": time.time(),
            **(metadata or {})
        }
        
        success_count = 0
        
        logger.info(f'[Session:{session_id}] 正在发送消息 (WebSocket={id(session.websocket)})')
        if await session.send_message(message):
            success_count = 1
            logger.info(f'[Session:{session_id}] 发送成功')
        else:
            logger.warning(f'[Session:{session_id}] 发送失败')
            # 清理失败的连接
            self.remove_session(session.websocket)
            
        logger.info(f'[Session:{session_id}] 广播原始消息完成: 成功{success_count}/总计1')
        return success_count
        
    async def broadcast_to_all(self, message_type: str, content: Any, 
                             source: str = "系统", metadata: Dict = None) -> int:
        """向所有WebSocket连接广播消息"""
        total_sent = 0
        with self._lock:
            session_ids = list(self._sessions.keys())
            
        for session_id in session_ids:
            sent = await self.broadcast_to_session(session_id, message_type, content, source, metadata)
            total_sent += sent
            
        logger.info(f'全局广播消息完成，总发送数: {total_sent}')
        return total_sent
        
    def get_session_count(self) -> int:
        """获取会话总数"""
        with self._lock:
            return len(self._sessions)
            
    def get_connection_count(self) -> int:
        """获取连接总数"""
        with self._lock:
            return len(self._websockets)
            
    def get_session_stats(self) -> Dict[str, Any]:
        """获取会话统计信息（单连接模式）"""
        with self._lock:
            stats = {
                "total_sessions": len(self._sessions),
                "total_connections": len(self._websockets),
                "session_details": {}
            }
            
            for session_id, session in self._sessions.items():
                stats["session_details"][session_id] = {
                    "connection_count": 1,  # 单连接模式固定为1
                    "connection": {
                        "created_at": session.created_at,
                        "last_ping": session.last_ping,
                        "is_alive": session.is_alive(),
                        "metadata": session.metadata
                    }
                }
                
        return stats
        
    async def _handle_ping(self, websocket: web.WebSocketResponse, data: Dict[str, Any]):
        """处理心跳消息"""
        session = self.get_session(websocket)
        if session:
            session.update_ping()
            await session.send_message({"type": "pong", "timestamp": time.time()})
            
    async def _handle_login(self, websocket: web.WebSocketResponse, data: Dict[str, Any]):
        """处理登录消息"""
        session_id = data.get('session_id', data.get('sessionid', str(int(time.time()))))
        
        # 确保session_id为字符串类型，避免类型不一致问题
        if isinstance(session_id, int):
            session_id = str(session_id)
        elif not isinstance(session_id, str):
            session_id = str(session_id)
        
        # 添加会话
        session = self.add_session(session_id, websocket)
        
        # 触发连接事件
        await self._emit_event('session_connected', session=session, data=data)
        
        # 发送登录确认
        await session.send_message({
            "type": "login_success",
            "data": {
                "session_id": session_id,
                "message": "WebSocket连接成功",
                "timestamp": time.time()
            }
        })
        
    async def handle_websocket_message(self, websocket: web.WebSocketResponse, data: Dict[str, Any]):
        """处理WebSocket消息"""
        message_type = data.get('type')
        
        if message_type in self._message_handlers:
            try:
                await self._message_handlers[message_type](websocket, data)
            except Exception as e:
                logger.error(f'消息处理器执行失败 {message_type}: {e}')
                session = self.get_session(websocket)
                if session:
                    await session.send_message({
                        "type": "error",
                        "data": {
                            "message": f"消息处理失败: {str(e)}",
                            "original_type": message_type
                        }
                    })
        else:
            logger.warning(f'未知消息类型: {message_type}')
            
    async def websocket_handler(self, request) -> web.WebSocketResponse:
        """WebSocket连接处理器"""
        ws = web.WebSocketResponse()
        await ws.prepare(request)
        
        logger.info('新的WebSocket连接建立')
        
        try:
            async for msg in ws:
                if msg.type == WSMsgType.TEXT:
                    try:
                        data = json.loads(msg.data)
                        await self.handle_websocket_message(ws, data)
                    except json.JSONDecodeError:
                        logger.error('收到无效的JSON数据')
                    except Exception as e:
                        logger.error(f'处理WebSocket消息时出错: {e}')
                        
                elif msg.type == WSMsgType.ERROR:
                    logger.error(f'WebSocket错误: {ws.exception()}')
                    break
                elif msg.type == WSMsgType.CLOSE:
                    logger.info('WebSocket连接正常关闭')
                    break
                    
        except ConnectionResetError:
            logger.warning('WebSocket连接被远程主机重置')
        except ConnectionAbortedError:
            logger.warning('WebSocket连接被中止')
        except Exception as e:
            logger.error(f'WebSocket连接错误: {e}')
        finally:
            # 清理会话
            session = self.remove_session(ws)
            if session:
                await self._emit_event('session_disconnected', session=session)
            logger.info('WebSocket连接已关闭')
            
        return ws
        
    def get_expired_sessions(self, timeout: int = 60) -> List[WebSocketSession]:
        """获取过期的会话列表"""
        current_time = time.time()
        expired_sessions = []
        
        with self._lock:
            for session in self._websockets.values():
                if current_time - session.last_ping > timeout:
                    expired_sessions.append(session)
                    
        return expired_sessions
        
    async def cleanup_dead_connections(self):
        """清理死连接"""
        dead_websockets = []
        
        with self._lock:
            for websocket, session in self._websockets.items():
                if not session.is_alive():
                    dead_websockets.append(websocket)
                    
        for websocket in dead_websockets:
            self.remove_session(websocket)
            
        if dead_websockets:
            logger.info(f'清理了 {len(dead_websockets)} 个死连接')
            
        return len(dead_websockets)


# 全局统一WebSocket管理器实例
_unified_manager = UnifiedWebSocketManager()


def get_unified_manager() -> UnifiedWebSocketManager:
    """获取统一WebSocket管理器实例"""
    return _unified_manager


# 兼容性接口，保持与原有代码的兼容
async def broadcast_message_to_session(session_id: str, message_type: str, content: str, 
                                     source: str = "数字人回复", model_info: str = None, 
                                     request_source: str = "页面"):
    """兼容性接口：向指定会话广播消息"""
    metadata = {}
    if model_info:
        metadata['model_info'] = model_info
    if request_source:
        metadata['request_source'] = request_source
        
    return await _unified_manager.broadcast_to_session(
        session_id, message_type, content, source, metadata
    )