"""
WebSocket 嵌套子任务测试
测试父 Agent 调用子 Agent 的场景
"""

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 NestedTaskTestClient:
    """嵌套任务测试客户端"""

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

    async def run_tests(self):
        """运行所有测试"""
        print("=" * 80)
        print("WebSocket 嵌套子任务测试")
        print("=" * 80)

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

            # 嵌套任务测试
            await self.test_nested_subtask()

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

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

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

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

    async def test_nested_subtask(self):
        """测试嵌套子任务"""
        print("\n[测试 2] 嵌套子任务测试")
        print("-" * 80)

        try:
            self.received_events = []
            subtask_events = {
                "started": [],
                "completed": [],
                "text_chunks": [],
                "tool_calls": []
            }

            # 注册事件处理器
            def on_subtask_started(event):
                self.received_events.append(("subtask_started", event))
                data = event.get("data", {})
                session_id = event.get("session_id", "")
                print(f"\n  📌 子任务开始:")
                print(f"     Session ID: {session_id}")
                print(f"     Parent ID: {data.get('parent_session_id')}")
                print(f"     描述: {data.get('description')}")
                print(f"     Agent 类型: {data.get('agent_type')}")
                subtask_events["started"].append(event)

            def on_subtask_completed(event):
                self.received_events.append(("subtask_completed", event))
                data = event.get("data", {})
                session_id = event.get("session_id", "")
                print(f"\n  ✅ 子任务完成:")
                print(f"     Session ID: {session_id}")
                print(f"     Parent ID: {data.get('parent_session_id')}")
                print(f"     描述: {data.get('description')}")
                print(f"     摘要: {data.get('summary', '')[:100]}...")
                subtask_events["completed"].append(event)

            def on_text_chunk(event):
                self.received_events.append(("text_chunk", event))
                data = event.get("data", {})
                session_id = event.get("session_id", "")
                content = data.get("content", "")
                if content.strip():
                    print(f"  💬 [{session_id[:8]}] {content[:50]}...")
                subtask_events["text_chunks"].append(event)

            def on_tool_call_start(event):
                self.received_events.append(("tool_call_start", event))
                data = event.get("data", {})
                session_id = event.get("session_id", "")
                print(f"  🔧 [{session_id[:8]}] 工具调用: {data.get('tool_name')}")
                subtask_events["tool_calls"].append(event)

            def on_tool_call_end(event):
                self.received_events.append(("tool_call_end", event))
                data = event.get("data", {})
                session_id = event.get("session_id", "")
                success = "✅" if data.get("success") else "❌"
                print(f"  {success} [{session_id[:8]}] 工具完成: {data.get('tool_name')}")

            def on_message_complete(event):
                self.received_events.append(("message_complete", event))
                session_id = event.get("session_id", "")
                print(f"\n  ✔️  消息完成: {session_id[:8]}")

            # 注册所有事件
            self.client.on("subtask_started", on_subtask_started)
            self.client.on("subtask_completed", on_subtask_completed)
            self.client.on("text_chunk", on_text_chunk)
            self.client.on("tool_call_start", on_tool_call_start)
            self.client.on("tool_call_end", on_tool_call_end)
            self.client.on("message_complete", on_message_complete)

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

            # 发送需要调用子 Agent 的消息
            print("\n发送消息: 请使用 task 工具创建一个子 Agent 来探索当前项目的结构")
            await self.client.send_message(
                "请使用 task 工具创建一个 explore 类型的子 Agent，"
                "让它探索当前项目的 src 目录结构，并列出主要的模块"
            )

            # 等待足够长的时间让子任务完成
            print("\n等待子任务执行...")
            await asyncio.sleep(30)

            # 取消接收任务
            receive_task.cancel()
            try:
                await receive_task
            except asyncio.CancelledError:
                pass

            # 验证结果
            print("\n" + "=" * 80)
            print("验证测试结果")
            print("=" * 80)

            # 1. 检查是否收到 subtask_started 事件
            if subtask_events["started"]:
                print("✅ 收到 subtask_started 事件")
                for event in subtask_events["started"]:
                    session_id = event.get("session_id", "")
                    # 验证 session_id 是层级的
                    if "." in session_id:
                        print(f"   ✅ Session ID 是层级的: {session_id}")
                    else:
                        print(f"   ❌ Session ID 不是层级的: {session_id}")
            else:
                print("❌ 未收到 subtask_started 事件")

            # 2. 检查是否收到 subtask_completed 事件
            if subtask_events["completed"]:
                print("✅ 收到 subtask_completed 事件")
            else:
                print("❌ 未收到 subtask_completed 事件")

            # 3. 检查是否收到子任务的文本输出
            if subtask_events["text_chunks"]:
                print(f"✅ 收到 {len(subtask_events['text_chunks'])} 个文本块")
            else:
                print("❌ 未收到子任务的文本输出")

            # 4. 检查是否收到工具调用事件
            if subtask_events["tool_calls"]:
                print(f"✅ 收到 {len(subtask_events['tool_calls'])} 个工具调用")
            else:
                print("⚠️  未收到工具调用事件（可能子任务没有使用工具）")

            # 5. 验证 session_id 层级关系
            print("\n验证 Session ID 层级关系:")
            parent_session_id = self.client.session_id
            print(f"  父 Session ID: {parent_session_id}")

            for event in subtask_events["started"]:
                sub_session_id = event.get("session_id", "")
                parent_id = event.get("data", {}).get("parent_session_id", "")
                print(f"  子 Session ID: {sub_session_id}")
                print(f"  声明的父 ID: {parent_id}")

                # 验证子 session_id 以父 session_id 开头
                if sub_session_id.startswith(parent_session_id + "."):
                    print(f"  ✅ 子 Session ID 正确继承父 ID")
                else:
                    print(f"  ❌ 子 Session ID 未正确继承父 ID")

            # 总体评估
            success = (
                len(subtask_events["started"]) > 0 and
                len(subtask_events["completed"]) > 0 and
                len(subtask_events["text_chunks"]) > 0
            )

            if success:
                print("\n✅ 嵌套子任务测试通过")
                self.test_results.append(("嵌套子任务", True, ""))
            else:
                error_msg = f"started={len(subtask_events['started'])}, completed={len(subtask_events['completed'])}, text={len(subtask_events['text_chunks'])}"
                print(f"\n❌ 嵌套子任务测试失败: {error_msg}")
                self.test_results.append(("嵌套子任务", False, error_msg))

        except Exception as e:
            print(f"❌ 嵌套子任务测试失败: {str(e)}")
            import traceback
            traceback.print_exc()
            self.test_results.append(("嵌套子任务", False, str(e)))

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

        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("-" * 80)
        print(f"总计: {passed}/{total} 通过")
        print("=" * 80)

        # 打印所有收到的事件统计
        print("\n事件统计:")
        event_counts = {}
        for event_type, _ in self.received_events:
            event_counts[event_type] = event_counts.get(event_type, 0) + 1

        for event_type, count in sorted(event_counts.items()):
            print(f"  {event_type}: {count}")


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

    # 检查是否在交互模式下运行
    import sys
    if sys.stdin.isatty():
        input("按 Enter 继续测试...")
    else:
        print("非交互模式，自动开始测试...")
        await asyncio.sleep(1)

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


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