import asyncio
import argparse
import json
import socket
import uuid
import websockets
import time
from typing import Dict, Any, Optional, Set
import logging

logger = logging.getLogger(__name__)

# Default browser function timeout in seconds
DEFAULT_BROWSER_FUNCTION_TIMEOUT = 180.0

# 允许连接的IP地址列表
ALLOWED_IPS = [
    "127.0.0.1",  # localhost
    "::1",  # localhost IPv6
]

class WebSocketManager:
    def __init__(self):
        self._connected_clients: Set[websockets.WebSocketServerProtocol] = set()
        self._client_connections: Dict[str, Dict[str, Any]] = {}
        self._server_started = False
        self._server = None
        self._response_futures: Dict[str, asyncio.Future] = {}
        self._port: int = 8765  # 默认端口，可通过 start_server(port=...) 覆盖
        self._shutting_down = False  # 标记是否正在关闭，用于快速退出等待循环

    def is_server_started(self):
        return self._server_started

    async def start_server(self, port: int = 8765):
        """启动WebSocket服务器
        
        Args:
            port: 监听端口，默认 8765。可通过 daemon 模式传入自定义端口。
        """
        # 防止重复启动：如果服务已在运行，直接返回
        if self._server_started:
            logger.info(f"WebSocket 服务器已在端口 {self._port} 上运行，跳过重复启动")
            return

        try:
            # 获取本机IP地址
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            try:
                s.connect(('8.8.8.8', 80))
                local_ip = s.getsockname()[0]
            except Exception:
                local_ip = '127.0.0.1'
            finally:
                s.close()
            
            # 确保本地IP在允许列表中
            if local_ip not in ALLOWED_IPS:
                ALLOWED_IPS.append(local_ip)
                logger.info(f"添加本地IP到允许列表: {local_ip}")
            
            # 启动服务器，只监听本地地址
            self._server = await websockets.serve(
                self._handle_client,
                "127.0.0.1",  # 只监听本地地址
                port,
                ping_interval=None,  # 禁用ping以减少网络开销
                max_size=10*1024*1024, # 10 MB
            )
            self._port = port
            self._server_started = True
            logger.info(f"WebSocket 服务器已启动在 ws://127.0.0.1:{port}")
            logger.info(f"允许连接的IP地址: {', '.join(ALLOWED_IPS)}")
            
        except Exception as e:
            logger.error(f"启动WebSocket服务器失败 (端口 {port}): {e}")
            raise

    async def _handle_client(self, client_ws: websockets.WebSocketServerProtocol):
        """处理WebSocket客户端连接

        Only one browser extension client is expected at a time.  When a new
        connection arrives we gracefully replace the previous one (if any),
        rather than blindly closing everything.
        """
        client_ip = client_ws.remote_address[0]

        # Replace the previous connection (only one browser client allowed)
        if self._connected_clients:
            old_clients = list(self._connected_clients)
            self._connected_clients.clear()
            self._client_connections.clear()
            for old_client in old_clients:
                try:
                    await old_client.close()
                    logger.info(f"已替换旧客户端连接 (IP: {old_client.remote_address[0]})")
                except Exception:
                    pass  # old connection may already be gone

        # Register the new client
        client_id = str(uuid.uuid4())
        self._client_connections[client_id] = {
            'ip': client_ip,
            'connected_at': asyncio.get_event_loop().time(),
            'ws': client_ws
        }
        self._connected_clients.add(client_ws)
        
        logger.info(f"新客户端连接 (ID: {client_id}, IP: {client_ip}), 当前连接数: {len(self._connected_clients)}")
        
        try:
            # 发送初始化消息，要求客户端确认连接
            init_message = {
                "type": "init",
                "client_id": client_id,
                "message": "请发送 connection_confirm 消息以确认连接"
            }
            logger.info(f"准备发送初始化消息给客户端 {client_id}: {init_message}")
            await client_ws.send(json.dumps(init_message))
            logger.info(f"已发送初始化消息给客户端 {client_id}")
            
            # 处理接收到的消息
            async for message in client_ws:
                try:
                    logger.info(f"收到客户端 {client_id} 的消息: {message}")
                    data = json.loads(message)
                    
                    # 处理连接确认消息
                    if data.get("type") == "connection_confirm":
                        logger.info(f"客户端 {client_id} 已确认连接")
                    
                    # 处理动作响应消息
                    elif "type" in data:
                        action = data.get("type")
                        if action in self._response_futures:
                            future = self._response_futures[action]
                            if not future.done():
                                # 直接返回字符串格式的结果
                                if action == "get_state":
                                    future.set_result(data)
                                elif action == "get_screenshot":
                                    future.set_result(data)                                    
                                elif action == "get_clickable_elements":
                                    future.set_result(data)
                                else:
                                    future.set_result(data)
                    
                    # 处理其他类型的响应消息
                    else:
                        # 获取消息中的第一个键作为动作类型
                        action = data['actionName']
                        if action in self._response_futures:
                            future = self._response_futures[action]
                            if not future.done():
                                future.set_result(data['actionResult'])
                                
                except json.JSONDecodeError:
                    logger.info(f"收到来自客户端 {client_id} 的原始消息: {message}")
                except Exception as e:
                    logger.error(f"处理客户端 {client_id} 消息时出错: {e}")
                    
        except websockets.exceptions.ConnectionClosed as e:
            logger.info(f"客户端 {client_id} 断开连接 - 状态码: {e.code}, 原因: {e.reason}")
            self._handle_connection_closed(client_ws, client_id)
        except Exception as e:
            logger.error(f"处理客户端 {client_id} 消息时出错: {e}")
            self._handle_connection_closed(client_ws, client_id)

    def _handle_connection_closed(self, client_ws: websockets.WebSocketServerProtocol, client_id: str):
        """处理连接关闭"""
        logger.info(f"客户端 {client_id} 断开连接")
        self._connected_clients.discard(client_ws)
        if client_id in self._client_connections:
            del self._client_connections[client_id]
        logger.info(f"当前连接数: {len(self._connected_clients)}")

    def _purge_dead_clients(self):
        """Remove clients whose underlying connection is already closed."""
        dead: set = set()
        for client in self._connected_clients:
            # websockets >= 14.x removed the `.closed` attribute;
            # use `.state` (an enum) instead.
            try:
                is_closed = client.state.name != "OPEN"
            except AttributeError:
                # Fallback for older versions that still have `.closed`
                is_closed = getattr(client, "closed", False)
            if is_closed:
                dead.add(client)
        for client in dead:
            self._connected_clients.discard(client)
            for cid, conn in list(self._client_connections.items()):
                if conn['ws'] == client:
                    del self._client_connections[cid]
                    logger.info(f"已清理断开的客户端 {cid}")
                    break

    async def send_message(self, message: Dict[str, Any], retry_times=1) -> str:
        """Send a message to the connected browser client and wait for a response.

        The method keeps the existing connection alive across calls.  It will
        only clean up connections that are *already dead*, never proactively
        close a healthy connection.
        """
        # Purge stale connections first
        self._purge_dead_clients()

        # Wait for a client to connect, but respect shutdown signal and timeout
        _wait_timeout = 60  # 最多等待60秒
        _wait_elapsed = 0
        while not self._connected_clients:
            if self._shutting_down:
                return "Error: Server is shutting down"
            if _wait_elapsed >= _wait_timeout:
                logger.warning(f"等待客户端连接超时（{_wait_timeout}秒），放弃本次请求")
                return f"Error: Waiting for client connection timed out after {_wait_timeout}s"
            logger.info(f"等待客户端连接... ({_wait_elapsed}/{_wait_timeout}s)")
            await asyncio.sleep(1)
            _wait_elapsed += 1
            
        # Determine action key from the message
        if "type" in message:
            action = message["type"]
        else:
            action = message.get("actionName")

        future = asyncio.Future()
        self._response_futures[action] = future
        
        try:
            logger.info(f"发送消息: {message}")
            tasks = []
            closed_clients: set = set()
            
            for client in self._connected_clients:
                try:
                    # Find the client_id for logging
                    client_id = None
                    for cid, conn in self._client_connections.items():
                        if conn['ws'] == client:
                            client_id = cid
                            break
                    
                    if client_id:
                        logger.info(f"向客户端 {client_id} 发送消息")
                        tasks.append(client.send(json.dumps(message)))
                    else:
                        logger.warning("找不到客户端ID，跳过发送消息")
                        closed_clients.add(client)
                except websockets.exceptions.ConnectionClosed:
                    logger.warning("客户端连接已关闭，跳过发送消息")
                    closed_clients.add(client)
                except Exception as e:
                    logger.error(f"检查客户端连接状态时出错: {e}")
                    closed_clients.add(client)
            
            # Clean up only the connections that are actually dead
            for client in closed_clients:
                self._connected_clients.discard(client)
                for cid, conn in list(self._client_connections.items()):
                    if conn['ws'] == client:
                        del self._client_connections[cid]
                        break
            
            if tasks:
                await asyncio.gather(*tasks)
                logger.info(f"消息已发送给 {len(tasks)} 个客户端")
            else:
                logger.warning("没有可用的客户端连接")
                return "Error: No active client connections"
            
            # Wait for the response
            try:
                response = await asyncio.wait_for(future, timeout=DEFAULT_BROWSER_FUNCTION_TIMEOUT)
                logger.info(f"收到响应: {response}")
                return response
            except asyncio.TimeoutError:
                return f"等待响应超时: {message}"
            finally:
                self._response_futures.pop(action, None)
                    
        except Exception as e:
            # Clean up only dead connections, do NOT close healthy ones
            self._purge_dead_clients()

            if retry_times <= 3:
                logger.error(f"发送消息时出错 ({e}), retry {retry_times}/3")
                # Brief pause before retry to give the connection a moment
                await asyncio.sleep(0.5)
                return await self.send_message(message, retry_times + 1)
            return f"Error: {str(e)}"

    async def cleanup(self):
        """清理资源"""
        logger.info(f"开始清理连接 - 当前连接数: {len(self._connected_clients)}")
        
        # 设置关闭标志，让 send_message 中的等待循环立即退出
        self._shutting_down = True
        
        # 取消所有等待中的 response futures，避免阻塞
        for action, future in list(self._response_futures.items()):
            if not future.done():
                future.cancel()
        self._response_futures.clear()
        
        # 关闭所有客户端连接
        for client_ws in list(self._connected_clients):
            try:
                await client_ws.close()
            except Exception as e:
                logger.error(f"关闭客户端连接时出错: {e}")
        
        self._connected_clients.clear()
        self._client_connections.clear()
        logger.info("所有客户端连接已关闭")
        
        # 关闭 WebSocket 服务器，释放端口
        if self._server:
            self._server.close()
            await self._server.wait_closed()
            self._server = None
        
        self._server_started = False
        logger.info("WebSocket 服务器已关闭，端口已释放")


async def test_client_connection():
    """测试客户端连接逻辑"""
    try:
        # 连接到WebSocket服务器
        async with websockets.connect("ws://127.0.0.1:8765") as websocket:
            logger.info("客户端已连接到服务器")
            
            # 接收初始化消息
            init_msg = await websocket.recv()
            logger.info(f"收到服务器消息: {init_msg}")
            
            # 解析并确认连接
            init_data = json.loads(init_msg)
            if init_data.get("type") == "init":
                confirm_msg = {
                    "type": "connection_confirm",
                    "client_id": init_data["client_id"],
                    "message": "连接已确认"
                }
                await websocket.send(json.dumps(confirm_msg))
                logger.info("已发送连接确认消息")
            
            # 测试发送一条消息
            test_msg = {
                "actionName": "test",
                "actionResult": "this is test message!",
            }
            await websocket.send(json.dumps(test_msg))
            logger.info("已发送测试消息")
            
    except Exception as e:
        logger.error(f"客户端测试失败: {e}")
        raise

async def main():
    """主入口函数"""
    parser = argparse.ArgumentParser(description='WebSocket服务器管理')
    parser.add_argument('--test', action='store_true', help='运行客户端测试')
    args = parser.parse_args()
    
    manager = WebSocketManager()
    
    try:
        # 启动服务器
        await manager.start_server()
        
        if args.test:
            # 运行客户端测试
            await test_client_connection()
        else:
            # 持续运行服务器
            logger.info("服务器运行中，按Ctrl+C退出...")
            while True:
                await asyncio.sleep(1)
                
    except KeyboardInterrupt:
        logger.info("收到中断信号，准备关闭...")
    except Exception as e:
        logger.error(f"运行出错: {e}")
    finally:
        # 清理资源
        await manager.cleanup()
        logger.info("服务器已关闭")

if __name__ == "__main__":
    # 直接运行时配置日志，确保 logger 输出不会丢失
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )
    asyncio.run(main())
