#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
WebSocket 实时推送验证脚本

功能:
1. 查询可用的测试用例
2. 创建执行任务
3. 连接 WebSocket 监听实时推送
4. 触发执行
5. 验证 WebSocket 消息是否正常接收

运行方式:
    python test_websocket_flow.py
"""

import asyncio
import json
import sys
import httpx
import websockets
from datetime import datetime
from pathlib import Path

# 添加项目路径
sys.path.insert(0, str(Path(__file__).parent.parent))

# API 配置
API_BASE_URL = "http://localhost:8001"
WS_BASE_URL = "ws://localhost:8001"


async def get_test_cases():
    """获取可用的测试用例"""
    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.get(f"{API_BASE_URL}/api/cases", params={"limit": 5})
        data = response.json()
        return data.get("items", [])


async def create_execution(case_ids: list):
    """创建执行任务"""
    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(
            f"{API_BASE_URL}/api/executions",
            json={
                "case_ids": case_ids,
                "name": f"WebSocket测试-{datetime.now().strftime('%H%M%S')}",
                "environment": "test",
                "config": {"headless": True, "auto_login": True}
            }
        )
        data = response.json()
        return data.get("id")


async def run_execution(execution_id: str):
    """触发执行"""
    async with httpx.AsyncClient(timeout=300.0) as client:
        response = await client.post(
            f"{API_BASE_URL}/api/executions/{execution_id}/run",
            json={"headless": True, "auto_login": True}
        )
        data = response.json()
        return data


async def listen_websocket(execution_id: str):
    """监听 WebSocket 推送"""
    uri = f"{WS_BASE_URL}/api/executions/ws/{execution_id}"

    messages_received = []

    print(f"\n[{datetime.now().strftime('%H:%M:%S')}] [WS] Connecting to: {uri}")

    try:
        async with websockets.connect(uri) as ws:
            print(f"[{datetime.now().strftime('%H:%M:%S')}] [WS] Connected!")

            # 发送心跳
            await ws.send(json.dumps({"type": "heartbeat"}))

            # 监听消息
            while True:
                try:
                    message = await asyncio.wait_for(ws.recv(), timeout=120.0)
                    data = json.loads(message)
                    msg_type = data.get("type", "unknown")

                    messages_received.append(data)

                    print(f"[{datetime.now().strftime('%H:%M:%S')}] [WS] Received: {msg_type}")

                    # 详细打印关键消息
                    if msg_type in ("execution_start", "case_start", "case_complete", "execution_complete"):
                        msg_data = data.get("data", {})
                        if msg_type == "execution_start":
                            print(f"    Execution: {msg_data.get('name')}, Total cases: {msg_data.get('total_cases')}")
                        elif msg_type == "case_start":
                            print(f"    Case: {msg_data.get('case_name')}")
                        elif msg_type == "case_complete":
                            print(f"    Case: {msg_data.get('case_name')}, Status: {msg_data.get('status')}")
                        elif msg_type == "execution_complete":
                            print(f"    Passed: {msg_data.get('passed')}, Failed: {msg_data.get('failed')}")
                            print(f"    Pass rate: {msg_data.get('pass_rate')}%")

                    # 执行完成，退出
                    if msg_type in ("execution_complete", "execution_cancelled"):
                        break

                except asyncio.TimeoutError:
                    print(f"[{datetime.now().strftime('%H:%M:%S')}] [WS] Timeout waiting for message")
                    break

    except Exception as e:
        print(f"[{datetime.now().strftime('%H:%M:%S')}] [WS] Error: {str(e)}")

    return messages_received


async def main():
    """主流程"""
    print("=" * 70)
    print("WebSocket Real-time Push Validation Test")
    print("=" * 70)

    # Step 1: 获取测试用例
    print(f"\n[{datetime.now().strftime('%H:%M:%S')}] [Step 1] Getting test cases...")
    cases = await get_test_cases()

    if not cases:
        print("Error: No test cases found!")
        return

    # 选择需要登录态的用例（auto_login=True）
    suitable_cases = [c for c in cases if c.get("config", {}).get("auto_login", False)]
    if not suitable_cases:
        suitable_cases = cases[:2]  # 任意选择前两个

    case_ids = [c["id"] for c in suitable_cases[:2]]
    print(f"  Selected cases: {[c['name'] for c in suitable_cases[:2]]}")
    print(f"  Case IDs: {case_ids}")

    # Step 2: 创建执行任务
    print(f"\n[{datetime.now().strftime('%H:%M:%S')}] [Step 2] Creating execution...")
    execution_id = await create_execution(case_ids)
    print(f"  Execution ID: {execution_id}")

    if not execution_id:
        print("Error: Failed to create execution!")
        return

    # Step 3: 启动 WebSocket 监听（后台任务）
    print(f"\n[{datetime.now().strftime('%H:%M:%S')}] [Step 3] Starting WebSocket listener...")

    websocket_task = asyncio.create_task(listen_websocket(execution_id))

    # 等待 WebSocket 连接建立
    await asyncio.sleep(2)

    # Step 4: 触发执行
    print(f"\n[{datetime.now().strftime('%H:%M:%S')}] [Step 4] Triggering execution...")
    result = await run_execution(execution_id)
    print(f"  Execution result: {result}")

    # Step 5: 等待 WebSocket 完成
    print(f"\n[{datetime.now().strftime('%H:%M:%S')}] [Step 5] Waiting for WebSocket messages...")
    messages = await websocket_task

    # 验证结果
    print("\n" + "=" * 70)
    print("Validation Results")
    print("=" * 70)
    print(f"  Total messages received: {len(messages)}")

    # 检查关键消息
    message_types = [m.get("type") for m in messages]

    expected_types = ["execution_start", "case_complete", "execution_complete"]
    found_types = [t for t in expected_types if t in message_types]

    print(f"  Expected message types: {expected_types}")
    print(f"  Found message types: {found_types}")

    # 验证是否所有关键消息都收到了
    if all(t in message_types for t in expected_types):
        print("\n[SUCCESS] WebSocket real-time push is working correctly!")
    elif len(messages) > 0:
        print(f"\n[PARTIAL] Received {len(messages)} messages, but some expected types missing")
        print(f"  Missing: {[t for t in expected_types if t not in message_types]}")
    else:
        print("\n[FAILED] No WebSocket messages received!")

    print("=" * 70)


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