#!/usr/bin/env python
# -*- coding: utf-8 """
WebSocket 客户端测试脚本

用于验证后端 WebSocket 实时推送功能是否正常工作。

运行方式:
    python test_websocket_client.py <execution_id>

示例:
    python test_websocket_client.py exec_abc123
"""

import asyncio
import json
import sys
import websockets
from datetime import datetime


async def test_websocket(execution_id: str):
    """测试 WebSocket 连接"""
    uri = f"ws://localhost:8001/api/executions/ws/{execution_id}"

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

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

            # 发送心跳消息
            await ws.send(json.dumps({"type": "heartbeat"}))
            print(f"[{datetime.now().strftime('%H:%M:%S')}] Sent heartbeat")

            # 监听消息（最多等待 60 秒）
            message_count = 0
            try:
                while True:
                    message = await asyncio.wait_for(ws.recv(), timeout=60.0)
                    message_count += 1

                    try:
                        data = json.loads(message)
                        msg_type = data.get("type", "unknown")
                        timestamp = data.get("timestamp", "")

                        print(f"[{datetime.now().strftime('%H:%M:%S')}] Received message #{message_count}:")
                        print(f"  Type: {msg_type}")
                        print(f"  Timestamp: {timestamp}")
                        print(f"  Data: {json.dumps(data.get('data', {}), indent=2, ensure_ascii=False)}")

                        # 如果收到执行完成消息，退出
                        if msg_type in ("execution_complete", "execution_cancelled"):
                            print(f"[{datetime.now().strftime('%H:%M:%S')}] Execution finished, closing connection...")
                            break

                    except json.JSONDecodeError:
                        print(f"[{datetime.now().strftime('%H:%M:%S')}] Raw message: {message[:200]}")

            except asyncio.TimeoutError:
                print(f"[{datetime.now().strftime('%H:%M:%S')}] Timeout (60s), closing connection...")

            print(f"[{datetime.now().strftime('%H:%M:%S')}] Total messages received: {message_count}")

    except websockets.exceptions.InvalidURI:
        print(f"Error: Invalid WebSocket URI: {uri}")
    except websockets.exceptions.InvalidHandshake:
        print(f"Error: WebSocket handshake failed")
    except Exception as e:
        print(f"Error: {str(e)}")


def main():
    """主函数"""
    if len(sys.argv) < 2:
        print("Usage: python test_websocket_client.py <execution_id>")
        print("Example: python test_websocket_client.py exec_abc123")
        sys.exit(1)

    execution_id = sys.argv[1]
    asyncio.run(test_websocket(execution_id))


if __name__ == "__main__":
    main()