import asyncio
import json
import httpx
import os

BASE_URL = "http://100.64.0.34:12000"       # dbg21

async def test_health():
    async with httpx.AsyncClient(base_url=BASE_URL) as client:
        response = await client.get("/health")
        if response.status_code != 200:
            print(f"Health Check Status Code: {response.status_code}")
            return
        print(f"Health Check: {response.json()}")

async def test_list_workflows():
    async with httpx.AsyncClient(base_url=BASE_URL, timeout=30.0) as client:
        response = await client.get("/workflows/available")
        if response.status_code != 200:
            print(f"List Available Workflows Status Code: {response.status_code}")
            return
        print(f"Available Workflows: {response.json()}")

async def test_register_models():
    async with httpx.AsyncClient(base_url=BASE_URL) as client:
        response = await client.post("/models/register", json={
            "profiles": [
                {
                    # Local used fields
                    "name": "DeepSeek Chat",
                    "model_id": "deepseek-chat",
                    "provider_name": "DeepSeek",
                    # Connection credentials
                    "model_name": "deepseek-chat",
                    "base_url": "https://api.deepseek.com/v1",
                    "api_key": os.getenv("DEEPSEEK_API_KEY"),
                }
            ]
        })
        if response.status_code != 200:
            print(f"Register Models Status Code: {response.status_code}")
            return
        print(f"Register Models: {response.json()}")

async def test_load_workflow():
    async with httpx.AsyncClient(base_url=BASE_URL, timeout=30.0) as client:
        response = await client.post("/workflows/load", json={
            "workflow_id": "alayamem_chat",
            "version": "1.0.0",
            "init_args": {
                "alayamem_url": "http://100.64.0.5:5556",
            },
        })
        if response.status_code != 200:
            print(f"Load Workflow Status Code: {response.status_code}")
            return
        print(f"Load Workflow: {response.json()}")

async def test_streaming_execution():
    """测试流式执行接口"""
    url = f"{BASE_URL}/workflows/execute/stream"
    
    payload = {
        "workflow_id": "alayamem_chat",
        "version": "1.0.0",
        "inputs": {
            "messages": [{"role": "user", "content": "你好"}]
        },
        "context": {
            "user_id": "user",
            "session_id": "test_session_multi_turn",
            "chat_model_id": "deepseek-chat",
        },
        "metadata": {
            "project": {
                "name": "test",
                "id": "cmkaw73cu000ap607qkp34k1i"
            },
            "org": {
                "name": "sustech",
                "id": "cmkaw5tyw0001p607w96xf0j0"
            }
        }
    }

    print("\n--- Starting Streaming Execution ---")
    async with httpx.AsyncClient(timeout=60.0) as client:
        # 使用 astream 发起流式请求
        async with client.stream("POST", url, json=payload) as response:
            if response.status_code != 200:
                print(f"Error: {response.status_code}, {await response.aread()}")
                return

            # 逐行读取响应 (NDJSON)
            async for line in response.aiter_lines():
                if not line:
                    continue
                
                try:
                    event_data = json.loads(line)
                    # 打印实时收到的事件
                    print(f"[Stream Event]: {event_data}")
                    
                    # 检查是否有错误事件
                    if event_data.get("event") == "error":
                        print(f"Error in stream: {event_data.get('error')}")
                        
                except json.JSONDecodeError:
                    print(f"[Raw Data]: {line}")

    print("--- Streaming Execution Finished ---\n")

async def main():
    # 1. 检查健康状态
    try:
        await test_health()
    except httpx.ConnectError:
        print("无法连接到服务器，请确保 server.py 正在运行 (uvicorn server:app --reload)")
        return

    # 2. 列出工作流 (alayahub 未正确配置的话，可能会比较长)
    await test_list_workflows()

    # 3. 注册模型
    await test_register_models()

    # 4. 加载工作流
    await test_load_workflow()

    # 5. 执行流式调用
    await test_streaming_execution()

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