"""
WebSocket监控模块
监控WebSocket连接和消息
"""
import json
import asyncio
from typing import Optional, List, Dict, Any, Callable
from playwright.async_api import Page, WebSocket
from loguru import logger
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum


class WSMessageType(Enum):
    """WebSocket消息类型"""
    TEXT = "text"
    BINARY = "binary"
    PING = "ping"
    PONG = "pong"
    CLOSE = "close"


@dataclass
class WSMessage:
    """WebSocket消息"""
    type: WSMessageType
    data: Any
    direction: str  # "sent" or "received"
    timestamp: datetime = field(default_factory=datetime.now)
    size: int = 0
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "type": self.type.value,
            "data": self.data,
            "direction": self.direction,
            "timestamp": self.timestamp.isoformat(),
            "size": self.size,
        }


@dataclass
class WSConnection:
    """WebSocket连接信息"""
    url: str
    opened_at: datetime
    closed_at: Optional[datetime] = None
    messages: List[WSMessage] = field(default_factory=list)
    is_open: bool = True
    
    @property
    def duration_ms(self) -> float:
        """连接持续时间"""
        end = self.closed_at or datetime.now()
        return (end - self.opened_at).total_seconds() * 1000
    
    @property
    def message_count(self) -> int:
        """消息数量"""
        return len(self.messages)
    
    @property
    def sent_count(self) -> int:
        """发送消息数量"""
        return sum(1 for m in self.messages if m.direction == "sent")
    
    @property
    def received_count(self) -> int:
        """接收消息数量"""
        return sum(1 for m in self.messages if m.direction == "received")


class WebSocketMonitor:
    """
    WebSocket监控器
    监控页面WebSocket连接和消息
    """
    
    def __init__(self, page: Page):
        self.page = page
        self._connections: Dict[str, WSConnection] = {}
        self._is_monitoring = False
        self._on_message: Optional[Callable[[str, WSMessage], None]] = None
        self._on_connect: Optional[Callable[[str], None]] = None
        self._on_disconnect: Optional[Callable[[str], None]] = None
    
    def start(
        self,
        on_message: Optional[Callable[[str, WSMessage], None]] = None,
        on_connect: Optional[Callable[[str], None]] = None,
        on_disconnect: Optional[Callable[[str], None]] = None,
    ):
        """
        开始监控
        
        Args:
            on_message: 消息回调
            on_connect: 连接回调
            on_disconnect: 断开回调
        """
        if self._is_monitoring:
            logger.warning("WebSocket监控已在运行")
            return
        
        self._is_monitoring = True
        self._on_message = on_message
        self._on_connect = on_connect
        self._on_disconnect = on_disconnect
        
        self.page.on("websocket", self._on_websocket)
        
        logger.info("WebSocket监控已启动")
    
    def stop(self):
        """停止监控"""
        if not self._is_monitoring:
            return
        
        self._is_monitoring = False
        self.page.remove_listener("websocket", self._on_websocket)
        
        logger.info("WebSocket监控已停止")
    
    def _on_websocket(self, ws: WebSocket):
        """WebSocket事件处理"""
        url = ws.url
        connection = WSConnection(
            url=url,
            opened_at=datetime.now(),
        )
        
        self._connections[url] = connection
        
        logger.info(f"WebSocket连接: {url}")
        
        if self._on_connect:
            self._on_connect(url)
        
        # 监听消息
        ws.on("framesreceived", lambda frames: self._on_frames_received(url, frames))
        ws.on("framessent", lambda frames: self._on_frames_sent(url, frames))
        ws.on("close", lambda: self._on_ws_close(url))
    
    def _on_frames_received(self, url: str, frames):
        """接收消息"""
        for frame in frames:
            message = self._create_message(frame, "received")
            self._connections[url].messages.append(message)
            
            logger.debug(f"WS接收: {url}, type={message.type.value}")
            
            if self._on_message:
                self._on_message(url, message)
    
    def _on_frames_sent(self, url: str, frames):
        """发送消息"""
        for frame in frames:
            message = self._create_message(frame, "sent")
            self._connections[url].messages.append(message)
            
            logger.debug(f"WS发送: {url}, type={message.type.value}")
            
            if self._on_message:
                self._on_message(url, message)
    
    def _on_ws_close(self, url: str):
        """连接关闭"""
        if url in self._connections:
            self._connections[url].is_open = False
            self._connections[url].closed_at = datetime.now()
            
            logger.info(f"WebSocket关闭: {url}")
            
            if self._on_disconnect:
                self._on_disconnect(url)
    
    def _create_message(self, frame, direction: str) -> WSMessage:
        """创建消息对象"""
        # 获取消息类型
        msg_type = WSMessageType.TEXT
        
        # 获取数据
        data = frame.payload
        if isinstance(data, bytes):
            msg_type = WSMessageType.BINARY
            data = data.hex()
        
        # 计算大小
        size = len(str(data))
        
        return WSMessage(
            type=msg_type,
            data=data,
            direction=direction,
            size=size,
        )
    
    def get_connections(self) -> List[WSConnection]:
        """获取所有连接"""
        return list(self._connections.values())
    
    def get_connection(self, url: str) -> Optional[WSConnection]:
        """获取指定连接"""
        return self._connections.get(url)
    
    def get_messages(self, url: Optional[str] = None) -> List[WSMessage]:
        """获取消息"""
        if url:
            conn = self._connections.get(url)
            return conn.messages if conn else []
        
        messages = []
        for conn in self._connections.values():
            messages.extend(conn.messages)
        return messages
    
    def get_json_messages(self, url: Optional[str] = None) -> List[Dict[str, Any]]:
        """获取JSON格式的消息"""
        messages = self.get_messages(url)
        json_messages = []
        
        for msg in messages:
            try:
                data = json.loads(msg.data)
                json_messages.append({
                    "message": msg.to_dict(),
                    "parsed": data,
                })
            except (json.JSONDecodeError, TypeError):
                pass
        
        return json_messages
    
    def clear(self):
        """清除记录"""
        self._connections.clear()
        logger.info("WebSocket记录已清除")
    
    def summary(self) -> Dict[str, Any]:
        """生成摘要"""
        connections = self.get_connections()
        
        total_messages = sum(c.message_count for c in connections)
        total_sent = sum(c.sent_count for c in connections)
        total_received = sum(c.received_count for c in connections)
        active = sum(1 for c in connections if c.is_open)
        
        return {
            "total_connections": len(connections),
            "active_connections": active,
            "total_messages": total_messages,
            "sent_messages": total_sent,
            "received_messages": total_received,
        }
    
    def print_summary(self):
        """打印摘要"""
        s = self.summary()
        
        print(f"\n=== WebSocket监控摘要 ===")
        print(f"连接总数: {s['total_connections']}")
        print(f"活跃连接: {s['active_connections']}")
        print(f"消息总数: {s['total_messages']}")
        print(f"发送: {s['sent_messages']}, 接收: {s['received_messages']}")


class WebSocketInterceptor:
    """
    WebSocket拦截器
    拦截和修改WebSocket消息
    """
    
    def __init__(self, page: Page):
        self.page = page
        self._interceptors: Dict[str, Callable] = {}
    
    async def intercept(
        self,
        url_pattern: str,
        on_received: Optional[Callable[[Any], Any]] = None,
        on_sent: Optional[Callable[[Any], Any]] = None,
    ):
        """
        拦截WebSocket消息
        
        Args:
            url_pattern: URL模式
            on_received: 接收消息处理函数
            on_sent: 发送消息处理函数
        """
        # 注入拦截脚本
        await self.page.evaluate(f"""
            () => {{
                const originalWebSocket = window.WebSocket;
                
                window.WebSocket = function(url, protocols) {{
                    const ws = new originalWebSocket(url, protocols);
                    
                    if ('{url_pattern}' === '*' || url.includes('{url_pattern}')) {{
                        // 拦截发送
                        const originalSend = ws.send;
                        ws.send = function(data) {{
                            console.log('[WS-INTERCEPT] Sent:', data);
                            originalSend.apply(ws, arguments);
                        }};
                        
                        // 拦截接收
                        ws.addEventListener('message', function(event) {{
                            console.log('[WS-INTERCEPT] Received:', event.data);
                        }});
                    }}
                    
                    return ws;
                }};
                
                window.WebSocket.prototype = originalWebSocket.prototype;
            }}
        """)
        
        self._interceptors[url_pattern] = {
            "on_received": on_received,
            "on_sent": on_sent,
        }
        
        logger.info(f"WebSocket拦截已设置: {url_pattern}")
    
    async def mock_response(
        self,
        url_pattern: str,
        messages: List[Dict[str, Any]],
        delay_ms: int = 0,
    ):
        """
        模拟WebSocket响应
        
        Args:
            url_pattern: URL模式
            messages: 消息列表
            delay_ms: 延迟
        """
        await self.page.evaluate(f"""
            () => {{
                const originalWebSocket = window.WebSocket;
                const mockMessages = {json.dumps(messages)};
                const delay = {delay_ms};
                
                window.WebSocket = function(url, protocols) {{
                    const ws = new originalWebSocket(url, protocols);
                    
                    if ('{url_pattern}' === '*' || url.includes('{url_pattern}')) {{
                        // 模拟消息
                        ws.addEventListener('open', function() {{
                            mockMessages.forEach((msg, index) => {{
                                setTimeout(() => {{
                                    ws.dispatchEvent(new MessageEvent('message', {{
                                        data: JSON.stringify(msg)
                                    }}));
                                }}, delay * index);
                            }});
                        }});
                    }}
                    
                    return ws;
                }};
                
                window.WebSocket.prototype = originalWebSocket.prototype;
            }}
        """)
        
        logger.info(f"WebSocket模拟已设置: {url_pattern}")


class WebSocketReplay:
    """
    WebSocket回放
    录制和回放WebSocket消息
    """
    
    def __init__(self, monitor: WebSocketMonitor):
        self.monitor = monitor
        self._recordings: Dict[str, List[WSMessage]] = {}
    
    def record(self, url: str):
        """开始录制"""
        self._recordings[url] = []
        
        def on_message(ws_url: str, message: WSMessage):
            if ws_url == url:
                self._recordings[url].append(message)
        
        self.monitor.start(on_message=on_message)
        logger.info(f"开始录制: {url}")
    
    def stop_recording(self, url: str) -> List[WSMessage]:
        """停止录制"""
        messages = self._recordings.get(url, [])
        logger.info(f"录制完成: {url}, {len(messages)} 条消息")
        return messages
    
    async def replay(self, page: Page, url: str, messages: List[WSMessage]):
        """回放消息"""
        # 连接到WebSocket
        await page.evaluate(f"""
            () => {{
                window.__replayWs = new WebSocket('{url}');
            }}
        """)
        
        # 回放消息
        for message in messages:
            if message.direction == "sent":
                await page.evaluate(f"""
                    () => {{
                        window.__replayWs.send('{message.data}');
                    }}
                """)
                await asyncio.sleep(0.1)  # 间隔
        
        logger.info(f"回放完成: {len(messages)} 条消息")
    
    def save_recording(self, url: str, filepath: str):
        """保存录制"""
        messages = self._recordings.get(url, [])
        
        data = {
            "url": url,
            "messages": [m.to_dict() for m in messages],
            "saved_at": datetime.now().isoformat(),
        }
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        logger.info(f"录制已保存: {filepath}")
    
    def load_recording(self, filepath: str) -> Dict[str, Any]:
        """加载录制"""
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        url = data["url"]
        self._recordings[url] = [
            WSMessage(
                type=WSMessageType(m["type"]),
                data=m["data"],
                direction=m["direction"],
                size=m["size"],
            )
            for m in data["messages"]
        ]
        
        logger.info(f"录制已加载: {filepath}")
        return data
