"""
WebSocket 测试脚本

测试用例：
1. 用户输入"你好"，验证流式返回
2. 用户输入文件分析请求，验证权限请求和授权
3. 用户输入网络搜索请求，验证权限请求和授权
"""

import asyncio
import json
import sys
from websockets import connect

WS_URL = "ws://localhost:8000/ws/test-session-001"


async def send_message(ws, content: str):
    """发送用户消息"""
    message = {
        "type": "user_message",
        "data": {"content": content}
    }
    await ws.send(json.dumps(message))
    print(f"\n>>> 发送: {content}")


async def send_permission_response(ws, request_id: str, approved: bool = True):
    """发送权限响应"""
    message = {
        "type": "permission_response",
        "data": {
            "request_id": request_id,
            "approved": approved,
            "remember": False
        }
    }
    await ws.send(json.dumps(message))
    print(f">>> 权限响应: approved={approved}")


async def receive_events(ws, timeout: float = 60.0, auto_approve: bool = True):
    """接收并处理事件"""
    text_buffer = []
    pending_permission_id = None

    try:
        while True:
            try:
                raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
                event = json.loads(raw)
                event_type = event.get("type", "")
                data = event.get("data", {})

                # 调试：打印所有收到的事件类型
                print(f"[DEBUG] 收到事件: {event_type}")

                if event_type == "connected":
                    print(f"<<< 已连接: session_id={data.get('session_id')}, agent={data.get('agent_type')}")

                elif event_type == "text_chunk":
                    chunk = data.get("content", "")
                    text_buffer.append(chunk)
                    print(chunk, end="", flush=True)

                elif event_type == "thinking_start":
                    print("<<< [思考中...]")

                elif event_type == "thinking_end":
                    pass

                elif event_type == "tool_call_start":
                    tool = data.get("tool", "")
                    args = data.get("args", {})
                    print(f"\n<<< [工具调用] {tool}: {json.dumps(args, ensure_ascii=False)[:100]}...")

                elif event_type == "tool_call_end":
                    tool = data.get("tool", "")
                    success = data.get("success", False)
                    result = data.get("result", "")[:200]
                    print(f"<<< [工具结果] {tool}: success={success}, result={result}...")

                elif event_type == "permission_request":
                    request_id = data.get("request_id", "")
                    permission = data.get("permission", "")
                    pattern = data.get("pattern", "")
                    print(f"\n<<< [权限请求] {permission}: {pattern}")

                    if auto_approve:
                        await send_permission_response(ws, request_id, approved=True)
                    else:
                        pending_permission_id = request_id

                elif event_type == "task_started":
                    print(f"<<< [任务开始] task_id={data.get('task_id')}")

                elif event_type == "task_completed":
                    print(f"\n<<< [任务完成]")
                    break

                elif event_type == "task_interrupted":
                    print(f"\n<<< [任务中断]")
                    break

                elif event_type == "task_error":
                    print(f"\n<<< [任务错误] {data.get('message')}")
                    break

                elif event_type == "error":
                    print(f"\n<<< [错误] {data.get('message')}")

                else:
                    print(f"<<< [未知事件] {event_type}: {data}")

            except asyncio.TimeoutError:
                print(f"\n<<< [超时] 等待响应超时 ({timeout}s)")
                break

    except Exception as e:
        print(f"\n<<< [异常] {e}")

    return "".join(text_buffer)


async def test_case_1():
    """测试用例 1: 用户输入"你好"，验证流式返回"""
    print("\n" + "="*60)
    print("测试用例 1: 用户输入'你好'，验证流式返回")
    print("="*60)

    async with connect(WS_URL) as ws:
        # 等待连接事件
        raw = await ws.recv()
        event = json.loads(raw)
        print(f"<<< 连接事件: {event}")

        # 发送消息
        await send_message(ws, "你好")

        # 接收响应
        response = await receive_events(ws, timeout=30.0)

        print("\n" + "-"*40)
        if response:
            print(f"✓ 测试通过: 收到流式响应 ({len(response)} 字符)")
            return True
        else:
            print("✗ 测试失败: 未收到响应")
            return False


async def test_case_2():
    """测试用例 2: 文件分析请求，验证权限请求和授权"""
    print("\n" + "="*60)
    print("测试用例 2: 文件分析请求，验证权限请求和授权")
    print("="*60)

    async with connect(WS_URL + "-case2") as ws:
        # 等待连接事件
        raw = await ws.recv()
        event = json.loads(raw)
        print(f"<<< 连接事件: {event}")

        # 发送消息
        await send_message(ws, "分析总结docs/CODE_AGENT_DESIGN.md，200字")

        # 接收响应（自动授权）
        response = await receive_events(ws, timeout=60.0, auto_approve=True)

        print("\n" + "-"*40)
        if response:
            print(f"✓ 测试通过: 收到分析结果 ({len(response)} 字符)")
            return True
        else:
            print("✗ 测试失败: 未收到响应")
            return False


async def test_case_3():
    """测试用例 3: 网络搜索请求，验证权限请求和授权"""
    print("\n" + "="*60)
    print("测试用例 3: 网络搜索请求，验证权限请求和授权")
    print("="*60)

    async with connect(WS_URL + "-case3") as ws:
        # 等待连接事件
        raw = await ws.recv()
        event = json.loads(raw)
        print(f"<<< 连接事件: {event}")

        # 发送消息
        await send_message(ws, "分析阿里巴巴最新财经新闻")

        # 接收响应（自动授权）
        response = await receive_events(ws, timeout=120.0, auto_approve=True)

        print("\n" + "-"*40)
        if response:
            print(f"✓ 测试通过: 收到搜索结果 ({len(response)} 字符)")
            return True
        else:
            print("✗ 测试失败: 未收到响应")
            return False


async def main():
    """运行所有测试"""
    if len(sys.argv) > 1:
        case = sys.argv[1]
        if case == "1":
            await test_case_1()
        elif case == "2":
            await test_case_2()
        elif case == "3":
            await test_case_3()
        else:
            print(f"未知测试用例: {case}")
    else:
        # 运行所有测试
        results = []

        results.append(("Case 1: 你好", await test_case_1()))
        await asyncio.sleep(2)

        results.append(("Case 2: 文件分析", await test_case_2()))
        await asyncio.sleep(2)

        results.append(("Case 3: 网络搜索", await test_case_3()))

        # 打印总结
        print("\n" + "="*60)
        print("测试总结")
        print("="*60)
        for name, passed in results:
            status = "✓ 通过" if passed else "✗ 失败"
            print(f"  {name}: {status}")


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