"""Python 3.11+ example. Install websockets, then:
python task_websocket.py http://localhost:8081 YOUR_API_KEY
Creates one TEXT task; task query recovers a lost/disconnected push.
Use HTTPS in production (the socket automatically uses WSS).
"""
import asyncio
import json
import sys
import time
import urllib.request
from urllib.parse import urlsplit, urlunsplit
import websockets

TERMINAL = {"SUCCEEDED", "FAILED", "EXPIRED"}

async def run(host, key):
    parts = urlsplit(host)
    url = urlunsplit(("wss" if parts.scheme == "https" else "ws", parts.netloc, "/socket/task.cgi", "", ""))
    def post(path, body):
        request = urllib.request.Request(host.rstrip("/") + path,
            json.dumps(body).encode(), {"Content-Type": "application/json", "apiKey": key})
        with urllib.request.urlopen(request, timeout=20) as response:
            payload = json.load(response)
        if payload["code"] != 0:
            raise RuntimeError(payload.get("msg") or str(payload["code"]))
        return payload
    updates = {}
    changed = asyncio.Event()
    disconnected = False
    async with websockets.connect(url, additional_headers={"apiKey": key}, open_timeout=10) as socket:
        ready = json.loads(await asyncio.wait_for(socket.recv(), 10))
        assert ready["type"] == "socket.ready"
        async def read():
            nonlocal disconnected
            try:
                async for frame in socket:
                    message = json.loads(frame)
                    if message["type"] == "task.updated":
                        task_id = message["task"]["taskId"]
                        if updates.get(task_id, {}).get("task", {}).get("state") not in TERMINAL:
                            updates[task_id] = message
                            changed.set()
            except websockets.ConnectionClosed:
                pass
            finally:
                disconnected = True
                changed.set()  # Wake the waiter immediately; query recovers a lost push.
        async def heartbeat():
            try:
                while True:
                    await asyncio.sleep(20)
                    await socket.send('{"type":"heartbeat"}')
            except websockets.ConnectionClosed:
                pass
        readers = [asyncio.create_task(read()), asyncio.create_task(heartbeat())]
        try:
            created = await asyncio.to_thread(post, "/api/validator/create.cgi",
                {"productCode": "TEXT", "adapter": {"comment": "What is 2 + 2?"}})
            task_id = created["data"]["taskId"]
            print(json.dumps(created))
            deadline = time.monotonic() + 1200
            while time.monotonic() < deadline:
                changed.clear()
                message = updates.get(task_id)
                if message and message["task"]["state"] in TERMINAL:
                    print(json.dumps(message)); return
                try:
                    # A pushed result wakes this immediately. Query only on lost connection/update.
                    await asyncio.wait_for(changed.wait(), 3 if disconnected else 30)
                except asyncio.TimeoutError:
                    result = await asyncio.to_thread(post, "/api/validator/query.cgi", {"taskId": task_id})
                    if result["data"]["state"] in TERMINAL:
                        print(json.dumps(result)); return
            raise TimeoutError(f"Query taskId {task_id} to recover the result")
        finally:
            for reader in readers: reader.cancel()
            await asyncio.gather(*readers, return_exceptions=True)

if __name__ == "__main__":
    asyncio.run(run(sys.argv[1], sys.argv[2]))