"""
完整的 WebSocket 测试
"""

import asyncio
import websockets
import json


async def test_full():
    """完整测试"""
    print("=" * 60)
    print("WebSocket 完整测试")
    print("=" * 60)

    try:
        async with websockets.connect("ws://localhost:8000/ws/test-full") as ws:
            print("\n✅ 连接成功")

            # 测试 1: 基础对话
            print("\n[测试 1] 基础对话")
            print("-" * 60)
            message = {
                "type": "user_message",
                "data": {"content": "你好"}
            }
            await ws.send(json.dumps(message))
            print("发送: 你好")

            # 接收所有事件
            events = []
            try:
                while True:
                    response = await asyncio.wait_for(ws.recv(), timeout=2.0)
                    event = json.loads(response)
                    event_type = event.get("type")
                    events.append(event_type)
                    print(f"  收到: {event_type}")

                    if event_type == "message_complete":
                        break
            except asyncio.TimeoutError:
                print("  (超时，停止接收)")

            if "text_chunk" in events and "message_complete" in events:
                print("✅ 基础对话测试通过")
            else:
                print(f"❌ 基础对话测试失败，收到事件: {events}")

            # 测试 2: 工具调用
            print("\n[测试 2] 工具调用")
            print("-" * 60)
            message = {
                "type": "user_message",
                "data": {"content": "北京天气怎么样"}
            }
            await ws.send(json.dumps(message))
            print("发送: 北京天气怎么样")

            events = []
            try:
                while True:
                    response = await asyncio.wait_for(ws.recv(), timeout=2.0)
                    event = json.loads(response)
                    event_type = event.get("type")
                    events.append(event_type)
                    print(f"  收到: {event_type}")

                    if event_type == "message_complete":
                        break
            except asyncio.TimeoutError:
                print("  (超时，停止接收)")

            if "tool_call_start" in events and "tool_call_end" in events:
                print("✅ 工具调用测试通过")
            else:
                print(f"❌ 工具调用测试失败，收到事件: {events}")

            # 测试 3: Agent 切换
            print("\n[测试 3] Agent 切换")
            print("-" * 60)
            message = {
                "type": "switch_agent",
                "data": {
                    "session_id": "test-full",
                    "agent_type": "code"
                }
            }
            await ws.send(json.dumps(message))
            print("发送: 切换到 code 模式")

            try:
                response = await asyncio.wait_for(ws.recv(), timeout=2.0)
                event = json.loads(response)
                if event.get("type") == "agent_switched":
                    print(f"  收到: agent_switched")
                    print("✅ Agent 切换测试通过")
                else:
                    print(f"❌ Agent 切换测试失败，收到: {event.get('type')}")
            except asyncio.TimeoutError:
                print("❌ Agent 切换测试失败：超时")

            # 测试 4: 命令执行
            print("\n[测试 4] 命令执行")
            print("-" * 60)
            message = {
                "type": "command",
                "data": {
                    "session_id": "test-full",
                    "command": "status",
                    "args": []
                }
            }
            await ws.send(json.dumps(message))
            print("发送: 执行 status 命令")

            try:
                response = await asyncio.wait_for(ws.recv(), timeout=2.0)
                event = json.loads(response)
                if event.get("type") == "command_result":
                    print(f"  收到: command_result")
                    print("✅ 命令执行测试通过")
                else:
                    print(f"❌ 命令执行测试失败，收到: {event.get('type')}")
            except asyncio.TimeoutError:
                print("❌ 命令执行测试失败：超时")

            print("\n" + "=" * 60)
            print("测试完成")
            print("=" * 60)

    except Exception as e:
        print(f"\n❌ 错误: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    asyncio.run(test_full())
