import asyncio
import time
from daystar_sdk.websocket_client import WebSocketRobotClient

"""
Task Management Demo
====================
演示 WebSocketRobotClient 任务相关接口使用：
1. create_task           创建持久化任务
2. execute_task          三种模式：仅脚本 / 脚本+task_id 调试 / 仅 task_id 执行已创建任务
3. task_status           查询任务状态（包含日志 / 错误信息）
4. pause_task / resume_task  暂停与恢复
5. stop_task             停止任务（可不传 task_id 停止当前）
6. get_all_tasks         获取全部任务ID列表
7. delete_task           删除任务

运行前注意：
- 需确保 WebSocket 服务端支持对应 /api/* 服务
- 根据实际主机/端口修改 client = WebSocketRobotClient(...)
- 这是一个串行演示，实际项目中可按需裁剪
"""

# 可选：状态枚举（假设后端 state 定义，可按实际返回值调整）
TASK_STATE_MAP = {
    0: "INIT",
    1: "RUNNING",
    2: "PAUSED",
    3: "FINISHED",
    4: "ERROR",
    5: "STOPPED"
}

async def poll_status(client: WebSocketRobotClient, task_id: str, timeout: float = 20.0, interval: float = 2.0):
    """轮询任务状态直到结束或超时"""
    start = time.time()
    while time.time() - start < timeout:
        ret = await client.task_status(task_id)
        if not ret.get("exists", True):
            print(f"[poll] 任务不存在 task_id={task_id}")
            return ret
        state = ret.get("state")
        state_name = TASK_STATE_MAP.get(state, str(state))
        print(f"[poll] state={state}({state_name}) message={ret.get('message','')} error={ret.get('error_message','')}")
        # 简单判断结束态（可按实际逻辑调整）
        if state_name in ("FINISHED", "ERROR", "STOPPED"):
            return ret
        await asyncio.sleep(interval)
    print("[poll] 超时未完成")
    return ret

async def main():
    client = WebSocketRobotClient(host="192.168.144.105", port=3232)
    await client.connect()

    # 1. 创建持久化任务
    create_ret = await client.create_task(
        task_id="demo_persist_task",
        script_content="import time\nfor i in range(5):\n    print('persist step', i); time.sleep(1)",
        timeout=120
    )
    print("create_task =>", create_ret)

    # 2. 一次性试运行（仅脚本）
    run_once_ret = await client.execute_task(
        script_content="print('single run task'); import math; print(math.sqrt(16))",
        timeout=15
    )
    print("execute_task (run once) =>", run_once_ret)

    # 3. 调试模式（task_id + script_content）
    debug_ret = await client.execute_task(
        task_id="demo_persist_task",
        script_content="print('debug run override code')",
        timeout=20
    )
    print("execute_task (debug) =>", debug_ret)

    # 4. 执行已存在的持久化任务（仅 task_id）
    run_persist_ret = await client.execute_task(
        task_id="demo_persist_task",
        timeout=60
    )
    print("execute_task (persist) =>", run_persist_ret)

    # 5. 轮询状态（可选）
    status_poll = await poll_status(client, "demo_persist_task", timeout=25, interval=3)
    print("poll final status =>", status_poll)

    # 6. 暂停任务
    pause_ret = await client.pause_task("demo_persist_task")
    print("pause_task =>", pause_ret)

    # 7. 恢复任务
    resume_ret = await client.resume_task("demo_persist_task")
    print("resume_task =>", resume_ret)

    # 8. 停止任务（指定ID）
    stop_ret = await client.stop_task("demo_persist_task")
    print("stop_task =>", stop_ret)

    # 9. 获取全部任务ID
    all_ret = await client.get_all_tasks()
    print("get_all_tasks =>", all_ret)

    # 10. 删除任务
    del_ret = await client.delete_task("demo_persist_task")
    print("delete_task =>", del_ret)

    await client.disconnect()

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