import asyncio, json, websockets


async def recv_response(ws, expected_id, timeout=10):
    """Read messages until we get the one matching our request id."""
    deadline = asyncio.get_event_loop().time() + timeout
    while True:
        remaining = deadline - asyncio.get_event_loop().time()
        if remaining <= 0:
            raise TimeoutError(f"No response for {expected_id} within {timeout}s")
        msg = json.loads(await asyncio.wait_for(ws.recv(), timeout=remaining))
        if msg.get("id") == expected_id:
            return msg
        # broadcast event — log and skip
        evt = msg.get("event", "unknown")
        data = json.dumps(msg.get("data", {}))[:200]
        print(f"  [event] {evt}: {data}")


async def drain_events(ws, seconds=4):
    """Drain broadcast events for N seconds."""
    try:
        while True:
            raw = await asyncio.wait_for(ws.recv(), timeout=seconds)
            msg = json.loads(raw)
            evt = msg.get("event", "?")
            data = json.dumps(msg.get("data", {}))[:200]
            print(f"  [event] {evt}: {data}")
    except asyncio.TimeoutError:
        pass


async def test_4c():
    async with websockets.connect("ws://127.0.0.1:6007") as ws:
        # 0. Stop any previously running game first
        print("--- stop (cleanup) ---")
        await ws.send(json.dumps({"id": "s0", "method": "stop", "params": {}}))
        s0 = await recv_response(ws, "s0")
        print(f"  status={s0.get('status')}")
        await drain_events(ws, 2)

        # 1. Run game — scene=null -> play main scene
        print("--- run ---")
        await ws.send(json.dumps({"id": "r1", "method": "run", "params": {}}))
        r1 = await recv_response(ws, "r1")
        print(f"  status={r1.get('status')}  result={r1.get('result')}")

        print("  waiting 6s for game to fully start...")
        await drain_events(ws, 6)

        # 1b. Check game status via editor_context
        print("--- editor_context ---")
        await ws.send(json.dumps({"id": "ec", "method": "editor_context", "params": {}}))
        ec = await recv_response(ws, "ec", timeout=5)
        game_running = ec.get("result", {}).get("game_running", "?")
        print(f"  game_running={game_running}")

        # 2. Test OLDER method: perf_snapshot
        print("--- perf_snapshot (older method) ---")
        await ws.send(json.dumps({"id": "p1", "method": "perf_snapshot", "params": {}}))
        try:
            p1 = await recv_response(ws, "p1", timeout=15)
            status_p = p1.get("status")
            body_p = json.dumps(p1.get("result", p1.get("error", {})), indent=2)[:400]
            print(f"  status={status_p}")
            print(f"  result={body_p}")
        except TimeoutError:
            print("  TIMEOUT!")

        # 3. Test NEW method: state_inspect
        print("--- state_inspect ---")
        await ws.send(json.dumps({"id": "r2", "method": "state_inspect", "params": {
            "queries": [
                {"autoload": "GameManager", "properties": ["current_day"]},
            ]
        }}))
        try:
            r2 = await recv_response(ws, "r2", timeout=15)
            status2 = r2.get("status")
            body2 = json.dumps(r2.get("result", r2.get("error", {})), indent=2)[:500]
            print(f"  status={status2}")
            print(f"  result={body2}")
        except TimeoutError:
            print("  TIMEOUT!")

        # 4. exec
        print("--- exec ---")
        await ws.send(json.dumps({"id": "r3", "method": "exec", "params": {
            "code": "func run():\n\treturn Engine.get_version_info()"
        }}))
        try:
            r3 = await recv_response(ws, "r3", timeout=15)
            status3 = r3.get("status")
            body3 = json.dumps(r3.get("result", r3.get("error", {})), indent=2)[:400]
            print(f"  status={status3}")
            print(f"  result={body3}")
        except TimeoutError:
            print("  TIMEOUT!")

        # 5. Stop
        print("--- stop ---")
        await ws.send(json.dumps({"id": "r5", "method": "stop", "params": {}}))
        r5 = await recv_response(ws, "r5")
        print(f"  status={r5.get('status')}")


asyncio.run(test_4c())
