"""xros_client.py — client library + demo for the XROS world-state service.

The demo proves DiOS's central claim: two independent "apps" (clients) observe
one shared world model. Client A spawns and moves an object; Client B, which
never touched that object, sees the spawn and the move arrive as events.

Run the server first (python xros_server.py), then:

    python xros_client.py --role spawner
    python xros_client.py --role observer

or run the self-contained demo (starts its own server in-process):

    python xros_client.py --demo
"""

from __future__ import annotations

import argparse
import asyncio
import json

import websockets


class WorldClient:
    """A minimal client for the world-state service."""

    def __init__(self, uri: str):
        self.uri = uri
        self.ws = None
        self.events: asyncio.Queue = asyncio.Queue()

    async def connect(self) -> None:
        self.ws = await websockets.connect(self.uri)
        asyncio.create_task(self._recv_loop())

    async def _recv_loop(self) -> None:
        async for raw in self.ws:
            self.events.put_nowait(json.loads(raw))

    async def call(self, msg: dict) -> dict:
        await self.ws.send(json.dumps(msg))
        # Read until the matching ack/error/snapshot/output for this call.
        while True:
            reply = await self.events.get()
            if reply.get("op") in ("ack", "error", "snapshot", "output"):
                return reply

    async def snapshot(self) -> list:
        return (await self.call({"op": "snapshot"}))["objects"]

    async def spawn(self, type: str, pose: dict | None = None, properties: dict | None = None) -> str:
        reply = await self.call({"op": "spawn", "type": type, "pose": pose, "properties": properties})
        return reply["id"]

    async def move(self, obj_id: str, pose: dict) -> None:
        await self.call({"op": "move", "id": obj_id, "pose": pose})

    async def command(self, text: str) -> str:
        """Run a spatial-shell command and return its text output."""
        reply = await self.call({"op": "command", "text": text})
        return reply.get("text", "")

    async def close(self) -> None:
        await self.ws.close()


async def run_spawner(uri: str) -> None:
    c = WorldClient(uri)
    await c.connect()
    obj_id = await c.spawn("cube", {"position": [0.0, 0.0, 0.0]}, {"color": "red"})
    print(f"[spawner] spawned {obj_id}")
    await asyncio.sleep(0.5)
    await c.move(obj_id, {"position": [1.0, 2.0, 0.0]})
    print(f"[spawner] moved {obj_id} to [1,2,0]")
    await asyncio.sleep(0.5)
    await c.close()


async def run_observer(uri: str) -> None:
    c = WorldClient(uri)
    await c.connect()
    print("[observer] connected; waiting for world events...")
    seen = 0
    try:
        while seen < 2:  # expect one spawn + one move
            event = await asyncio.wait_for(c.events.get(), timeout=5)
            if event.get("op") in ("spawn", "move"):
                seen += 1
                print(f"[observer] saw {event['op']}: {event.get('object', {}).get('id', '?')}")
    except asyncio.TimeoutError:
        print("[observer] timed out waiting for events")
    await c.close()


async def demo() -> None:
    """Self-contained proof: one in-process server, two clients, one shared world."""
    from xros_server import handle
    from xros_world import WorldState

    world = WorldState()

    async def server():
        async with websockets.serve(lambda ws: handle(ws, world), "127.0.0.1", 8765):
            await asyncio.Future()

    srv = asyncio.create_task(server())
    await asyncio.sleep(0.2)  # let the server bind

    uri = "ws://127.0.0.1:8765"
    spawner = asyncio.create_task(run_spawner(uri))
    observer = asyncio.create_task(run_observer(uri))
    await asyncio.gather(spawner, observer)

    print("\n--- final shared world state ---")
    for o in await world.snapshot():
        print(f"  {o['type']} {o['id'][:8]} at {o['pose']['position']}")

    srv.cancel()


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--uri", default="ws://127.0.0.1:8765")
    p.add_argument("--role", choices=["spawner", "observer"])
    p.add_argument("--demo", action="store_true")
    args = p.parse_args()

    if args.demo:
        asyncio.run(demo())
    elif args.role == "spawner":
        asyncio.run(run_spawner(args.uri))
    elif args.role == "observer":
        asyncio.run(run_observer(args.uri))
    else:
        p.print_help()