"""
WebSocket 协议测试脚本
"""

import asyncio
import sys
import os

# 添加项目路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

from client.websocket_client import WebSocketClient


class TestClient:
    """测试客户端"""

    def __init__(self, url: str = "ws://localhost:8000"):
        self.url = url
        self.client = None
        self.test_results = []

    async def run_tests(self):
        """运行所有测试"""
        print("=" * 60)
        print("WebSocket 协议测试")
        print("=" * 60)

        try:
            # 连接测试
            await self.test_connection()

            # 基础对话测试
            await self.test_basic_conversation()

            # Agent 切换测试
            await self.test_agent_switch()

            # 命令执行测试
            await self.test_command_execution()

            # 工具调用测试
            await self.test_tool_call()

            # 打印测试结果
            self.print_results()

        except Exception as e:
            print(f"\n❌ 测试失败: {str(e)}")
            import traceback
            traceback.print_exc()

    async def test_connection(self):
        """测试连接"""
        print("\n[测试 1] 连接测试")
        print("-" * 60)

        try:
            self.client = WebSocketClient(self.url, "test-session-1")
            await self.client.connect()
            print("✅ 连接成功")
            self.test_results.append(("连接测试", True, ""))
        except Exception as e:
            print(f"❌ 连接失败: {str(e)}")
            self.test_results.append(("连接测试", False, str(e)))
            raise

    async def test_basic_conversation(self):
        """测试基础对话"""
        print("\n[测试 2] 基础对话测试")
        print("-" * 60)

        try:
            # 注册事件处理器
            received_events = []

            def on_event(event):
                event_type = event.get("type")
                received_events.append(event_type)
                print(f"  收到事件: {event_type}")

            self.client.on("text_chunk", on_event)
            self.client.on("message_complete", on_event)

            # 启动接收循环
            receive_task = asyncio.create_task(self.client.receive_loop())

            # 发送消息
            print("发送消息: 你好")
            await self.client.send_message("你好")

            # 等待响应
            await asyncio.sleep(3)

            # 检查结果
            if "text_chunk" in received_events and "message_complete" in received_events:
                print("✅ 基础对话测试通过")
                self.test_results.append(("基础对话", True, ""))
            else:
                print(f"❌ 未收到预期事件，收到: {received_events}")
                self.test_results.append(("基础对话", False, f"收到事件: {received_events}"))

            receive_task.cancel()

        except Exception as e:
            print(f"❌ 基础对话测试失败: {str(e)}")
            self.test_results.append(("基础对话", False, str(e)))

    async def test_agent_switch(self):
        """测试 Agent 切换"""
        print("\n[测试 3] Agent 切换测试")
        print("-" * 60)

        try:
            received_events = []

            def on_agent_switched(event):
                received_events.append(event)
                data = event.get("data", {})
                print(f"  Agent 切换: {data.get('from_agent')} -> {data.get('to_agent')}")

            self.client.on("agent_switched", on_agent_switched)

            # 启动接收循环
            receive_task = asyncio.create_task(self.client.receive_loop())

            # 切换 Agent
            print("切换到 code 模式")
            await self.client.switch_agent("code")

            # 等待响应
            await asyncio.sleep(2)

            # 检查结果
            if received_events:
                print("✅ Agent 切换测试通过")
                self.test_results.append(("Agent切换", True, ""))
            else:
                print("❌ 未收到 agent_switched 事件")
                self.test_results.append(("Agent切换", False, "未收到事件"))

            receive_task.cancel()

        except Exception as e:
            print(f"❌ Agent 切换测试失败: {str(e)}")
            self.test_results.append(("Agent切换", False, str(e)))

    async def test_command_execution(self):
        """测试命令执行"""
        print("\n[测试 4] 命令执行测试")
        print("-" * 60)

        try:
            received_events = []

            def on_command_result(event):
                received_events.append(event)
                data = event.get("data", {})
                print(f"  命令结果: {data.get('command')} - {data.get('success')}")

            self.client.on("command_result", on_command_result)

            # 启动接收循环
            receive_task = asyncio.create_task(self.client.receive_loop())

            # 执行命令
            print("执行命令: /status")
            await self.client.execute_command("status")

            # 等待响应
            await asyncio.sleep(2)

            # 检查结果
            if received_events:
                print("✅ 命令执行测试通过")
                self.test_results.append(("命令执行", True, ""))
            else:
                print("❌ 未收到 command_result 事件")
                self.test_results.append(("命令执行", False, "未收到事件"))

            receive_task.cancel()

        except Exception as e:
            print(f"❌ 命令执行测试失败: {str(e)}")
            self.test_results.append(("命令执行", False, str(e)))

    async def test_tool_call(self):
        """测试工具调用"""
        print("\n[测试 5] 工具调用测试")
        print("-" * 60)

        try:
            received_events = []

            def on_tool_start(event):
                received_events.append("tool_call_start")
                data = event.get("data", {})
                print(f"  工具调用开始: {data.get('tool_name')}")

            def on_tool_end(event):
                received_events.append("tool_call_end")
                data = event.get("data", {})
                print(f"  工具调用结束: {data.get('tool_name')} - {data.get('success')}")

            self.client.on("tool_call_start", on_tool_start)
            self.client.on("tool_call_end", on_tool_end)

            # 启动接收循环
            receive_task = asyncio.create_task(self.client.receive_loop())

            # 发送需要工具调用的消息
            print("发送消息: 北京天气怎么样")
            await self.client.send_message("北京天气怎么样")

            # 等待响应
            await asyncio.sleep(10)

            # 检查结果
            if "tool_call_start" in received_events and "tool_call_end" in received_events:
                print("✅ 工具调用测试通过")
                self.test_results.append(("工具调用", True, ""))
            else:
                print(f"❌ 未收到预期事件，收到: {received_events}")
                self.test_results.append(("工具调用", False, f"收到事件: {received_events}"))

            receive_task.cancel()

        except Exception as e:
            print(f"❌ 工具调用测试失败: {str(e)}")
            self.test_results.append(("工具调用", False, str(e)))

    def print_results(self):
        """打印测试结果"""
        print("\n" + "=" * 60)
        print("测试结果汇总")
        print("=" * 60)

        passed = sum(1 for _, success, _ in self.test_results if success)
        total = len(self.test_results)

        for name, success, error in self.test_results:
            status = "✅ 通过" if success else "❌ 失败"
            print(f"{name:20s} {status}")
            if error:
                print(f"  错误: {error}")

        print("-" * 60)
        print(f"总计: {passed}/{total} 通过")
        print("=" * 60)


async def main():
    """主函数"""
    # 检查服务器是否运行
    print("请确保 WebSocket 服务器正在运行:")
    print("  python -m src.server.app")
    print()
    input("按 Enter 继续测试...")

    # 运行测试
    test_client = TestClient()
    await test_client.run_tests()


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