#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
WebSocket服务演示和测试
启动一个完整的WebSocket服务器，模拟客户端连接，测试内置函数功能
"""

import asyncio
import json
import threading
import time
import sys
import os
from datetime import datetime
from typing import Dict, Set

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

# 导入FastAPI和相关模块
try:
    import uvicorn
    from fastapi import FastAPI, WebSocket, WebSocketDisconnect
    from fastapi.middleware.cors import CORSMiddleware
    import websockets
except ImportError as e:
    print(f"需要安装依赖包: {e}")
    print("请运行: pip install fastapi uvicorn websockets")
    sys.exit(1)

# 导入WebSocket模块
try:
    from services.online_script.script_vwed_objects import create_vwed_object
    from utils.logger import get_logger
except ImportError as e:
    print(f"导入项目模块失败: {e}")
    print("请确保从项目根目录运行此脚本")
    sys.exit(1)

logger = get_logger("tests.websocket_service_demo")


class WebSocketTestServer:
    """WebSocket测试服务器"""
    
    def __init__(self, port: int = 8899):
        self.port = port
        self.app = FastAPI(title="WebSocket测试服务器")
        self.connections: Dict[str, WebSocket] = {}
        self.client_ips: Set[str] = set()
        
        # 添加CORS支持
        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )
        
        self.setup_routes()
    
    def setup_routes(self):
        """设置路由"""
        
        @self.app.websocket("/ws/{client_id}")
        async def websocket_endpoint(websocket: WebSocket, client_id: str):
            await websocket.accept()
            
            # 获取客户端IP
            client_ip = websocket.client.host if hasattr(websocket, 'client') else "unknown"
            self.connections[client_id] = websocket
            self.client_ips.add(client_ip)
            
            logger.info(f"客户端连接: {client_id} (IP: {client_ip})")
            
            try:
                # 发送欢迎消息
                await websocket.send_text(json.dumps({
                    "type": "welcome",
                    "message": f"欢迎 {client_id}",
                    "timestamp": datetime.now().isoformat()
                }))
                
                # 保持连接并监听消息
                while True:
                    try:
                        data = await websocket.receive_text()
                        logger.info(f"收到来自 {client_id} 的消息: {data}")
                        
                        # 回显消息
                        await websocket.send_text(json.dumps({
                            "type": "echo",
                            "original_message": data,
                            "timestamp": datetime.now().isoformat()
                        }))
                        
                    except WebSocketDisconnect:
                        break
                        
            except Exception as e:
                logger.error(f"WebSocket处理错误: {e}")
            finally:
                # 清理连接
                if client_id in self.connections:
                    del self.connections[client_id]
                logger.info(f"客户端断开连接: {client_id}")
        
        @self.app.get("/status")
        async def get_status():
            """获取服务器状态"""
            return {
                "active_connections": len(self.connections),
                "client_ids": list(self.connections.keys()),
                "client_ips": list(self.client_ips),
                "timestamp": datetime.now().isoformat()
            }
    
    def start_server(self):
        """启动服务器"""
        logger.info(f"启动WebSocket测试服务器，端口: {self.port}")
        
        # 在单独线程中运行服务器
        def run_server():
            uvicorn.run(self.app, host="0.0.0.0", port=self.port, log_level="warning")
        
        server_thread = threading.Thread(target=run_server, daemon=True)
        server_thread.start()
        
        # 等待服务器启动
        time.sleep(2)
        logger.info("WebSocket测试服务器已启动")
        
        return server_thread


class WebSocketTestClient:
    """WebSocket测试客户端"""
    
    def __init__(self, client_id: str, server_url: str = "ws://localhost:8899"):
        self.client_id = client_id
        self.server_url = f"{server_url}/ws/{client_id}"
        self.websocket = None
        self.messages = []
        self.connected = False
    
    async def connect(self):
        """连接到服务器"""
        try:
            self.websocket = await websockets.connect(self.server_url)
            self.connected = True
            logger.info(f"客户端 {self.client_id} 已连接到服务器")
            
            # 启动消息监听
            asyncio.create_task(self.listen_messages())
            
        except Exception as e:
            logger.error(f"客户端 {self.client_id} 连接失败: {e}")
            raise
    
    async def listen_messages(self):
        """监听服务器消息"""
        try:
            async for message in self.websocket:
                data = json.loads(message)
                self.messages.append(data)
                logger.info(f"客户端 {self.client_id} 收到消息: {data}")
                
        except Exception as e:
            logger.error(f"客户端 {self.client_id} 消息监听错误: {e}")
        finally:
            self.connected = False
    
    async def send_message(self, message: str):
        """发送消息到服务器"""
        if self.websocket and self.connected:
            await self.websocket.send(message)
            logger.info(f"客户端 {self.client_id} 发送消息: {message}")
    
    async def disconnect(self):
        """断开连接"""
        if self.websocket:
            await self.websocket.close()
            self.connected = False
            logger.info(f"客户端 {self.client_id} 已断开连接")


async def test_websocket_builtin_functions():
    """测试WebSocket内置函数"""
    logger.info("=== 测试WebSocket内置函数 ===")
    
    try:
        # 创建VWED对象
        vwed = create_vwed_object("test_script_websocket")
        
        # 等待客户端连接稳定
        await asyncio.sleep(1)
        
        # 测试获取客户端信息
        logger.info("--- 测试获取客户端信息 ---")
        try:
            client_ips = vwed.websocket.get_websocket_client_ip()
            client_names = vwed.websocket.get_websocket_client_name()
            
            logger.info(f"获取到的客户端IP: {client_ips}")
            logger.info(f"获取到的客户端名称: {client_names}")
        except Exception as e:
            logger.warning(f"获取客户端信息失败 (可能是因为使用了不同的连接管理器): {e}")
        
        # 测试发送消息功能
        logger.info("--- 测试发送消息功能 ---")
        
        # 注意：由于我们使用的是独立的测试服务器，这些测试可能会失败
        # 但可以展示如何使用这些函数
        
        try:
            # 测试根据IP发送消息
            vwed.websocket.send_msg_to_wsc_by_client_ip("来自VWED脚本的消息", "127.0.0.1")
            logger.info("尝试根据IP发送消息")
        except Exception as e:
            logger.info(f"根据IP发送消息测试 (预期可能失败): {e}")
        
        try:
            # 测试根据客户端名称发送消息
            vwed.websocket.send_msg_to_wsc_by_client_name("来自VWED脚本的消息", "test_client_1")
            logger.info("尝试根据客户端名称发送消息")
        except Exception as e:
            logger.info(f"根据客户端名称发送消息测试 (预期可能失败): {e}")
        
        logger.info("WebSocket内置函数测试完成")
        
    except Exception as e:
        logger.error(f"WebSocket内置函数测试失败: {e}")


async def run_websocket_demo():
    """运行WebSocket演示"""
    logger.info("开始WebSocket服务演示")
    
    # 启动测试服务器
    server = WebSocketTestServer()
    server_thread = server.start_server()
    
    try:
        # 创建测试客户端
        clients = []
        client_tasks = []
        
        # 创建多个客户端
        for i in range(3):
            client_id = f"test_client_{i+1}"
            client = WebSocketTestClient(client_id)
            clients.append(client)
            
            # 连接客户端
            await client.connect()
            await asyncio.sleep(0.5)  # 错开连接时间
        
        # 客户端发送测试消息
        for i, client in enumerate(clients):
            await client.send_message(f"测试消息来自客户端 {i+1}")
            await asyncio.sleep(0.2)
        
        # 等待消息处理
        await asyncio.sleep(2)
        
        # 测试WebSocket内置函数
        await test_websocket_builtin_functions()
        
        # 展示客户端收到的消息
        logger.info("=== 客户端消息汇总 ===")
        for client in clients:
            logger.info(f"客户端 {client.client_id} 收到 {len(client.messages)} 条消息")
            for msg in client.messages:
                logger.info(f"  - {msg}")
        
        # 断开客户端连接
        for client in clients:
            await client.disconnect()
        
        logger.info("✅ WebSocket服务演示完成")
        
    except Exception as e:
        logger.error(f"❌ WebSocket服务演示失败: {e}")
        raise
    
    finally:
        logger.info("清理测试环境")


def main():
    """主函数"""
    logger.info("WebSocket服务演示和测试开始")
    
    try:
        # 运行异步演示
        asyncio.run(run_websocket_demo())
        
    except KeyboardInterrupt:
        logger.info("用户中断演示")
    except Exception as e:
        logger.error(f"演示过程中发生错误: {e}")
        raise
    
    logger.info("WebSocket服务演示和测试结束")


if __name__ == "__main__":
    main()