import asyncio
import json
import sys

import websockets

WS_URL = sys.argv[1] if len(sys.argv) > 1 else "ws://127.0.0.1:8010/ws/generate"


async def main() -> None:
    payload = {
        "action": "generate",
        "payload": {
            "prompt": "a lonely man walking in a rainy neon street, cinematic, handheld camera",
            "negative_prompt": "blurry, deformed face, extra limbs, flicker",
            "quality_mode": "preview",
            "duration_sec": 1,
            "width": 320,
            "height": 240,
            "fps": 8,
            "steps": 8,
            "seed": 123456,
        },
    }

    async with websockets.connect(WS_URL, max_size=2**22) as ws:
        await ws.send(json.dumps(payload, ensure_ascii=False))
        while True:
            msg = await ws.recv()
            data = json.loads(msg)
            print(json.dumps(data, ensure_ascii=False))
            if data.get("event") == "result":
                break


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