#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：manager.py
模块描述：WebSocket 连接管理器，负责实时推送测试执行状态

作者：czj
创建日期：2026-07-09
最后修改：2026-07-09
"""

import json
import logging
from typing import Dict, Set, Optional, Any
from datetime import datetime

from fastapi import WebSocket, WebSocketDisconnect

logger = logging.getLogger(__name__)


class ConnectionManager:
    """
    WebSocket 连接管理器

    管理所有 WebSocket 连接，支持按执行ID分组广播消息。
    每个执行记录对应一个分组，订阅该执行的客户端会收到实时状态推送。

    Attributes:
        connections (Dict[str, Set[WebSocket]]): 按执行ID分组的连接集合

    Example:
        >>> manager = ConnectionManager()
        >>> # 在路由中接受连接
        >>> await manager.connect(websocket, "exec_001")
        >>> # 广播消息
        >>> await manager.broadcast("exec_001", {"type": "step_update", "data": {...}})
        >>> # 断开连接
        >>> await manager.disconnect(websocket, "exec_001")
    """

    def __init__(self):
        """初始化连接管理器"""
        # 按执行ID分组的连接集合
        self.connections: Dict[str, Set[WebSocket]] = {}
        # 反向映射：WebSocket -> 它订阅的执行ID列表
        self._ws_to_executions: Dict[WebSocket, Set[str]] = {}

    async def connect(self, websocket: WebSocket, execution_id: str) -> None:
        """
        接受 WebSocket 连接并注册到指定执行分组

        Args:
            websocket (WebSocket): WebSocket 连接实例
            execution_id (str): 执行记录ID

        Example:
            >>> @app.websocket("/ws/execution/{execution_id}")
            >>> async def websocket_endpoint(websocket: WebSocket, execution_id: str):
            >>>     await manager.connect(websocket, execution_id)
            >>>     try:
            >>>         while True:
            >>>             await websocket.receive_text()
            >>>     except WebSocketDisconnect:
            >>>         await manager.disconnect(websocket, execution_id)
        """
        await websocket.accept()

        # 注册到执行分组
        if execution_id not in self.connections:
            self.connections[execution_id] = set()
        self.connections[execution_id].add(websocket)

        # 记录反向映射
        if websocket not in self._ws_to_executions:
            self._ws_to_executions[websocket] = set()
        self._ws_to_executions[websocket].add(execution_id)

        logger.info(f"WebSocket 连接已建立: execution_id={execution_id}, 当前连接数={len(self.connections[execution_id])}")

    async def disconnect(self, websocket: WebSocket, execution_id: Optional[str] = None) -> None:
        """
        断开 WebSocket 连接

        如果指定了 execution_id，只从该分组移除；
        否则从所有关联的分组移除。

        Args:
            websocket (WebSocket): WebSocket 连接实例
            execution_id (Optional[str]): 执行记录ID
        """
        if execution_id:
            # 从指定分组移除
            if execution_id in self.connections:
                self.connections[execution_id].discard(websocket)
                if not self.connections[execution_id]:
                    del self.connections[execution_id]
        else:
            # 从所有关联分组移除
            exec_ids = self._ws_to_executions.get(websocket, set())
            for eid in exec_ids:
                if eid in self.connections:
                    self.connections[eid].discard(websocket)
                    if not self.connections[eid]:
                        del self.connections[eid]

        # 清理反向映射
        self._ws_to_executions.pop(websocket, None)

        execution_info = execution_id or "所有"
        logger.info(f"WebSocket 连接已断开: execution_id={execution_info}")

    async def broadcast(
        self,
        execution_id: str,
        message: Dict[str, Any],
        exclude: Optional[WebSocket] = None,
    ) -> None:
        """
        向指定执行分组的客户端广播消息

        Args:
            execution_id (str): 执行记录ID
            message (dict): 消息内容
            exclude (Optional[WebSocket]): 排除的客户端（不需要通知的）

        Example:
            >>> await manager.broadcast("exec_001", {
            ...     "type": "step_update",
            ...     "data": {
            ...         "execution_id": "exec_001",
            ...         "case_name": "正常登录测试",
            ...         "step_order": 1,
            ...         "status": "passed"
            ...     }
            ... })
        """
        if execution_id not in self.connections:
            return

        # 添加时间戳到消息
        if "timestamp" not in message:
            message["timestamp"] = datetime.now().isoformat()

        message_str = json.dumps(message, ensure_ascii=False)

        # 遍历该分组的连接并发送
        disconnected = []
        for websocket in self.connections[execution_id]:
            if websocket == exclude:
                continue

            try:
                await websocket.send_text(message_str)
            except WebSocketDisconnect:
                disconnected.append(websocket)
            except Exception as e:
                logger.error(f"WebSocket 发送失败: {str(e)}")
                disconnected.append(websocket)

        # 清理已断开的连接
        for ws in disconnected:
            await self.disconnect(ws, execution_id)

    async def send_to(
        self,
        websocket: WebSocket,
        message: Dict[str, Any],
    ) -> None:
        """
        向单个 WebSocket 客户端发送消息

        Args:
            websocket (WebSocket): 目标客户端
            message (dict): 消息内容
        """
        try:
            message_str = json.dumps(message, ensure_ascii=False)
            await websocket.send_text(message_str)
        except Exception as e:
            logger.error(f"发送消息失败: {str(e)}")

    def get_connection_count(self, execution_id: Optional[str] = None) -> int:
        """
        获取连接数统计

        Args:
            execution_id (Optional[str]): 指定执行ID，为None时返回总数

        Returns:
            int: 连接数
        """
        if execution_id:
            return len(self.connections.get(execution_id, set()))
        return sum(len(ws) for ws in self.connections.values())

    def get_active_executions(self) -> list:
        """
        获取有活跃连接的执行ID列表

        Returns:
            list: 活跃执行ID列表
        """
        return [
            {"execution_id": eid, "connections": len(ws_set)}
            for eid, ws_set in self.connections.items()
            if ws_set
        ]


# 全局 WebSocket 连接管理器实例
manager = ConnectionManager()