"""
简化的 WebSocket 嵌套测试 - 使用 Mock 验证事件流
"""

import asyncio
import json
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

from server.events import ProtocolEvent, EventType
from core.session_id import SessionIDGenerator


def test_session_id_generator():
    """测试 SessionIDGenerator"""
    print("=" * 80)
    print("测试 SessionIDGenerator")
    print("=" * 80)

    # 生成根 ID
    root_id = SessionIDGenerator.generate()
    print(f"根 ID: {root_id}")
    assert len(root_id) == 8
    assert SessionIDGenerator.validate(root_id)

    # 生成子 ID
    sub_id = SessionIDGenerator.create_subtask_id(root_id)
    print(f"子 ID: {sub_id}")
    assert sub_id.startswith(root_id + ".")
    assert SessionIDGenerator.get_depth(sub_id) == 1
    assert SessionIDGenerator.get_root_id(sub_id) == root_id

    # 生成孙 ID
    sub_sub_id = SessionIDGenerator.create_subtask_id(sub_id)
    print(f"孙 ID: {sub_sub_id}")
    assert sub_sub_id.startswith(sub_id + ".")
    assert SessionIDGenerator.get_depth(sub_sub_id) == 2
    assert SessionIDGenerator.get_root_id(sub_sub_id) == root_id

    print("✅ SessionIDGenerator 测试通过\n")
    return root_id, sub_id, sub_sub_id


def test_event_creation():
    """测试事件创建"""
    print("=" * 80)
    print("测试事件创建")
    print("=" * 80)

    root_id = "a3b5c7d9"
    sub_id = "a3b5c7d9.x1y2z3w4"

    # 测试 subtask_started 事件
    started_event = ProtocolEvent(
        EventType.SUBTASK_STARTED,
        root_id,  # 父 Agent 发送，使用父 session_id
        {
            "subtask_session_id": sub_id,
            "description": "测试子任务",
            "subagent_type": "explore"
        }
    )

    event_dict = started_event.to_dict()
    print(f"subtask_started 事件:")
    print(f"  type: {event_dict['type']}")
    print(f"  session_id: {event_dict['session_id']}")
    print(f"  subtask_session_id: {event_dict['data']['subtask_session_id']}")

    assert event_dict['type'] == 'subtask_started'
    assert event_dict['session_id'] == root_id
    assert event_dict['data']['subtask_session_id'] == sub_id

    # 测试子 Agent 的 text_chunk 事件
    text_event = ProtocolEvent(
        EventType.TEXT_CHUNK,
        sub_id,  # 子 Agent 发送，使用子 session_id
        {
            "content": "子 Agent 输出",
            "is_final": False
        }
    )

    event_dict = text_event.to_dict()
    print(f"\ntext_chunk 事件（子 Agent）:")
    print(f"  type: {event_dict['type']}")
    print(f"  session_id: {event_dict['session_id']}")
    print(f"  content: {event_dict['data']['content']}")

    assert event_dict['type'] == 'text_chunk'
    assert event_dict['session_id'] == sub_id

    # 测试 subtask_completed 事件
    completed_event = ProtocolEvent(
        EventType.SUBTASK_COMPLETED,
        root_id,  # 父 Agent 发送，使用父 session_id
        {
            "subtask_session_id": sub_id,
            "title": "测试子任务",
            "success": True,
            "duration_ms": 1000
        }
    )

    event_dict = completed_event.to_dict()
    print(f"\nsubtask_completed 事件:")
    print(f"  type: {event_dict['type']}")
    print(f"  session_id: {event_dict['session_id']}")
    print(f"  subtask_session_id: {event_dict['data']['subtask_session_id']}")
    print(f"  success: {event_dict['data']['success']}")

    assert event_dict['type'] == 'subtask_completed'
    assert event_dict['session_id'] == root_id
    assert event_dict['data']['subtask_session_id'] == sub_id

    print("\n✅ 事件创建测试通过\n")


def test_event_routing_logic():
    """测试事件路由逻辑"""
    print("=" * 80)
    print("测试事件路由逻辑")
    print("=" * 80)

    # 模拟 WebSocket 连接映射
    websocket_connections = {}

    root_id = "a3b5c7d9"
    sub_id = "a3b5c7d9.x1y2z3w4"
    sub_sub_id = "a3b5c7d9.x1y2z3w4.m5n6p7q8"

    # 注册连接（使用根 ID）
    websocket_connections[root_id] = "MockWebSocket"

    # 测试路由逻辑
    def get_websocket_for_session(session_id):
        """根据 session_id 获取 WebSocket 连接"""
        root = SessionIDGenerator.get_root_id(session_id)
        return websocket_connections.get(root)

    # 测试父 Agent 事件路由
    ws = get_websocket_for_session(root_id)
    print(f"父 Agent ({root_id}) -> WebSocket: {ws}")
    assert ws == "MockWebSocket"

    # 测试子 Agent 事件路由
    ws = get_websocket_for_session(sub_id)
    print(f"子 Agent ({sub_id}) -> WebSocket: {ws}")
    assert ws == "MockWebSocket"

    # 测试孙 Agent 事件路由
    ws = get_websocket_for_session(sub_sub_id)
    print(f"孙 Agent ({sub_sub_id}) -> WebSocket: {ws}")
    assert ws == "MockWebSocket"

    print("\n✅ 事件路由逻辑测试通过\n")


def test_complete_event_flow():
    """测试完整的事件流"""
    print("=" * 80)
    print("测试完整的事件流")
    print("=" * 80)

    root_id = "a3b5c7d9"
    sub_id = "a3b5c7d9.x1y2z3w4"

    events = []

    # 1. 父 Agent 输出
    events.append({
        "type": "text_chunk",
        "session_id": root_id,
        "data": {"content": "我来帮你探索..."}
    })

    # 2. 父 Agent 调用 task 工具
    events.append({
        "type": "tool_call_start",
        "session_id": root_id,
        "data": {"tool_name": "task"}
    })

    # 3. 子任务开始
    events.append({
        "type": "subtask_started",
        "session_id": root_id,
        "data": {"subtask_session_id": sub_id}
    })

    # 4. 子 Agent 输出
    events.append({
        "type": "text_chunk",
        "session_id": sub_id,
        "data": {"content": "正在搜索..."}
    })

    # 5. 子 Agent 完成
    events.append({
        "type": "message_complete",
        "session_id": sub_id,
        "data": {}
    })

    # 6. 子任务完成
    events.append({
        "type": "subtask_completed",
        "session_id": root_id,
        "data": {"subtask_session_id": sub_id, "success": True}
    })

    # 7. 父 Agent 继续
    events.append({
        "type": "text_chunk",
        "session_id": root_id,
        "data": {"content": "根据探索结果..."}
    })

    # 8. 父 Agent 完成
    events.append({
        "type": "message_complete",
        "session_id": root_id,
        "data": {}
    })

    print("完整事件流:")
    for i, event in enumerate(events, 1):
        session_id = event['session_id']
        depth = SessionIDGenerator.get_depth(session_id)
        indent = "  " * depth
        print(f"{i}. {indent}{event['type']} (session_id: {session_id})")

    # 验证事件顺序
    assert events[0]['type'] == 'text_chunk'
    assert events[2]['type'] == 'subtask_started'
    assert events[3]['session_id'] == sub_id
    assert events[5]['type'] == 'subtask_completed'

    print("\n✅ 完整事件流测试通过\n")


def main():
    """运行所有测试"""
    print("\n")
    print("=" * 80)
    print("WebSocket 嵌套子 Agent 事件流测试（Mock）")
    print("=" * 80)
    print("\n")

    try:
        test_session_id_generator()
        test_event_creation()
        test_event_routing_logic()
        test_complete_event_flow()

        print("=" * 80)
        print("✅ 所有测试通过！")
        print("=" * 80)
        print("\n")
        print("说明:")
        print("- SessionIDGenerator 正常工作")
        print("- 事件创建和序列化正常")
        print("- 事件路由逻辑正确")
        print("- 完整事件流符合预期")
        print("\n")
        print("下一步:")
        print("1. 启动 WebSocket 服务器: python start_server.py")
        print("2. 启动客户端: python start_client.py")
        print("3. 测试真实的 Agent 调用 task 工具")
        print("\n")

    except AssertionError as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return 1
    except Exception as e:
        print(f"\n❌ 错误: {e}")
        import traceback
        traceback.print_exc()
        return 1

    return 0


if __name__ == "__main__":
    sys.exit(main())
