"""
WebSocket 客户端测试脚本
用于测试 WebSocket 实时数据推送功能
"""
import asyncio
import websockets
import json
from typing import Optional


class WebSocketTestClient:
    """WebSocket 测试客户端"""

    def __init__(self, url: str = "ws://localhost:8000/ws"):
        """
        初始化测试客户端

        - **url**: WebSocket 服务器 URL
        """
        self.url = url
        self.websocket: Optional[websockets.WebSocketClientProtocol] = None
        self.connected = False

    async def connect(self, token: Optional[str] = None):
        """
        连接到 WebSocket 服务器

        - **token**: 可选的 JWT Token
        """
        uri = f"{self.url}"
        if token:
            uri += f"?token={token}"

        try:
            self.websocket = await websockets.connect(uri)
            self.connected = True
            print(f"[+] Connected to {self.url}")

            # 接收欢迎消息
            welcome_msg = await self.websocket.recv()
            print(f"[+] Welcome message: {welcome_msg}")
            return json.loads(welcome_msg)
        except Exception as e:
            print(f"[-] Failed to connect: {e}")
            raise

    async def disconnect(self):
        """断开连接"""
        if self.websocket:
            await self.websocket.close()
            self.connected = False
            print("[+] Disconnected from server")

    async def subscribe(self, subscription_types: list):
        """
        订阅数据类型

        - **subscription_types**: 订阅类型列表，如 ["tick", "order", "trade"]
        """
        message = {
            "type": "subscribe",
            "data": {
                "subscription": subscription_types
            }
        }
        await self.send(message)
        print(f"[+] Subscribed to: {subscription_types}")

    async def unsubscribe(self, subscription_types: list):
        """
        取消订阅

        - **subscription_types**: 订阅类型列表
        """
        message = {
            "type": "unsubscribe",
            "data": {
                "subscription": subscription_types
            }
        }
        await self.send(message)
        print(f"[+] Unsubscribed from: {subscription_types}")

    async def subscribe_symbol(self, symbols: list):
        """
        订阅品种行情

        - **symbols**: 品种代码列表，如 ["IF2501.CFFEX", "IH2501.CFFEX"]
        """
        message = {
            "type": "subscribe_symbol",
            "data": {
                "symbol": symbols
            }
        }
        await self.send(message)
        print(f"[+] Subscribed to symbols: {symbols}")

    async def unsubscribe_symbol(self, symbols: list):
        """
        取消品种订阅

        - **symbols**: 品种代码列表
        """
        message = {
            "type": "unsubscribe_symbol",
            "data": {
                "symbol": symbols
            }
        }
        await self.send(message)
        print(f"[+] Unsubscribed from symbols: {symbols}")

    async def send_ping(self):
        """发送心跳"""
        message = {
            "type": "ping",
            "data": {}
        }
        await self.send(message)
        print("[+] Ping sent")

    async def send(self, message: dict):
        """
        发送消息

        - **message**: 消息字典
        """
        if not self.websocket or not self.connected:
            raise Exception("Not connected to WebSocket server")

        await self.websocket.send(json.dumps(message))

    async def receive(self, timeout: Optional[float] = None):
        """
        接收消息

        - **timeout**: 超时时间（秒）
        """
        if not self.websocket or not self.connected:
            raise Exception("Not connected to WebSocket server")

        try:
            message = await asyncio.wait_for(self.websocket.recv(), timeout=timeout)
            return json.loads(message)
        except asyncio.TimeoutError:
            return None

    async def listen(self, duration: int = 10, print_messages: bool = True):
        """
        监听消息

        - **duration**: 监听时长（秒）
        - **print_messages**: 是否打印消息
        """
        print(f"\n[*] Listening for messages ({duration}s)...")
        messages = []

        try:
            while True:
                message = await asyncio.wait_for(self.websocket.recv(), timeout=duration)
                data = json.loads(message)
                messages.append(data)

                if print_messages:
                    msg_type = data.get("type", "unknown")
                    print(f"[*] Received {msg_type}: {json.dumps(data, ensure_ascii=False)[:200]}...")

        except asyncio.TimeoutError:
            print(f"\n[+] Listening finished. Received {len(messages)} messages")
            return messages


async def test_basic_connection():
    """测试基本连接功能"""
    print("\n" + "="*50)
    print("TEST 1: Basic Connection")
    print("="*50)

    client = WebSocketTestClient()

    try:
        # 连接
        await client.connect()
        await asyncio.sleep(1)

        # 测试心跳
        await client.send_ping()
        response = await client.receive(timeout=2)
        if response and response.get("type") == "pong":
            print("[+] Ping/Pong test passed")

        # 断开
        await client.disconnect()
        print("[+] Test passed")
        return True

    except Exception as e:
        print(f"[-] Test failed: {e}")
        return False


async def test_subscription():
    """测试订阅功能"""
    print("\n" + "="*50)
    print("TEST 2: Subscription")
    print("="*50)

    client = WebSocketTestClient()

    try:
        # 连接
        await client.connect()

        # 订阅多个数据类型
        await client.subscribe(["tick", "order", "trade", "position", "account", "log"])
        await asyncio.sleep(0.5)

        # 订阅品种
        await client.subscribe_symbol(["IF2501.CFFEX", "IH2501.CFFEX"])
        await asyncio.sleep(0.5)

        # 取消订阅
        await client.unsubscribe(["log"])
        await asyncio.sleep(0.5)

        # 断开
        await client.disconnect()
        print("[+] Test passed")
        return True

    except Exception as e:
        print(f"[-] Test failed: {e}")
        return False


async def test_message_reception():
    """测试消息接收功能"""
    print("\n" + "="*50)
    print("TEST 3: Message Reception")
    print("="*50)

    client = WebSocketTestClient()

    try:
        # 连接
        await client.connect()

        # 订阅所有类型
        await client.subscribe(["tick", "order", "trade", "position", "account", "log", "contract"])

        # 监听消息
        messages = await client.listen(duration=5)

        print(f"\n[+] Received messages by type:")
        msg_types = {}
        for msg in messages:
            msg_type = msg.get("type", "unknown")
            msg_types[msg_type] = msg_types.get(msg_type, 0) + 1

        for msg_type, count in msg_types.items():
            print(f"  - {msg_type}: {count}")

        # 断开
        await client.disconnect()
        print("[+] Test passed")
        return True

    except Exception as e:
        print(f"[-] Test failed: {e}")
        return False


async def test_heartbeat():
    """测试心跳机制"""
    print("\n" + "="*50)
    print("TEST 4: Heartbeat")
    print("="*50)

    client = WebSocketTestClient()

    try:
        # 连接
        await client.connect()

        # 发送多次心跳
        for i in range(5):
            await client.send_ping()
            response = await client.receive(timeout=2)
            if response and response.get("type") == "pong":
                print(f"[+] Heartbeat {i+1}/5 successful")
            await asyncio.sleep(1)

        # 断开
        await client.disconnect()
        print("[+] Test passed")
        return True

    except Exception as e:
        print(f"[-] Test failed: {e}")
        return False


async def test_error_handling():
    """测试错误处理"""
    print("\n" + "="*50)
    print("TEST 5: Error Handling")
    print("="*50)

    client = WebSocketTestClient()

    try:
        # 连接
        await client.connect()

        # 发送未知消息类型
        await client.send({
            "type": "unknown_type",
            "data": {}
        })

        response = await client.receive(timeout=2)
        if response and response.get("type") == "error":
            print("[+] Error response received correctly")

        # 断开
        await client.disconnect()
        print("[+] Test passed")
        return True

    except Exception as e:
        print(f"[-] Test failed: {e}")
        return False


async def test_authenticated_connection():
    """测试认证连接"""
    print("\n" + "="*50)
    print("TEST 6: Authenticated Connection")
    print("="*50)

    # 注意：需要有效的 JWT Token
    # 这里测试无效 Token 的情况
    client = WebSocketTestClient()

    try:
        # 使用无效 Token 连接
        await client.connect(token="invalid_token")
        print("[-] Should have failed with invalid token")
        await client.disconnect()
        return False

    except Exception as e:
        print(f"[+] Correctly rejected invalid token: {e}")
        return True


async def run_all_tests():
    """运行所有测试"""
    print("\n" + "="*50)
    print("WebSocket Test Suite")
    print("="*50)

    tests = [
        ("Basic Connection", test_basic_connection),
        ("Subscription", test_subscription),
        ("Message Reception", test_message_reception),
        ("Heartbeat", test_heartbeat),
        ("Error Handling", test_error_handling),
        ("Authenticated Connection", test_authenticated_connection),
    ]

    results = []

    for name, test_func in tests:
        try:
            result = await test_func()
            results.append((name, result))
        except Exception as e:
            print(f"[-] Test '{name}' crashed: {e}")
            results.append((name, False))

        await asyncio.sleep(1)

    # 汇总结果
    print("\n" + "="*50)
    print("Test Results Summary")
    print("="*50)

    passed = sum(1 for _, result in results if result)
    total = len(results)

    for name, result in results:
        status = "✓ PASS" if result else "✗ FAIL"
        print(f"{status}: {name}")

    print(f"\nTotal: {passed}/{total} tests passed")

    return passed == total


def main():
    """主函数"""
    # 运行测试
    success = asyncio.run(run_all_tests())

    if success:
        print("\n[+] All tests passed!")
        exit(0)
    else:
        print("\n[-] Some tests failed")
        exit(1)


if __name__ == "__main__":
    main()