"""xros_server.py — WebSocket server exposing the shared world-state API.

Runs a single WorldState and lets any number of clients read and mutate it
over a JSON-over-WebSocket protocol. This is the "world as a system service"
primitive from DiOS: apps consume one shared world model instead of each
reconstructing their own.

Protocol (JSON messages, one object per line):

  client -> server:
    {"op": "snapshot"}                          -> full current world
    {"op": "spawn", "type": "...", "pose": {...}, "properties": {...}}
    {"op": "move",  "id": "...", "pose": {...}}
    {"op": "remove", "id": "..."}

  server -> client (after any mutation, to every connected client):
    {"op": "spawn",  "object": {...}}
    {"op": "move",   "object": {...}}
    {"op": "remove", "id": "..."}

Run:  python xros_server.py --host 127.0.0.1 --port 8765
"""

from __future__ import annotations

import argparse
import asyncio
import json

import websockets

from xros_world import Pose, WorldState
from xros_shell import SpatialShell


async def handle(ws, world: WorldState, shell: SpatialShell) -> None:
    """One connection = one app. Subscribe it to the event bus, then serve it."""
    queue = world.subscribe()
    # Forward world events to this client.
    forward = asyncio.create_task(_forward(queue, ws))
    try:
        async for raw in ws:
            try:
                msg = json.loads(raw)
            except json.JSONDecodeError:
                await ws.send(json.dumps({"op": "error", "error": "bad json"}))
                continue
            await _dispatch(ws, world, shell, msg)
    finally:
        forward.cancel()
        world.unsubscribe(queue)


async def _forward(queue: asyncio.Queue, ws) -> None:
    while True:
        event = await queue.get()
        await ws.send(json.dumps(event))


async def _dispatch(ws, world: WorldState, shell: SpatialShell, msg: dict) -> None:
    op = msg.get("op")
    try:
        if op == "command":
            output = await shell.execute(msg.get("text", ""))
            await ws.send(json.dumps({"op": "output", "text": output}))
        elif op == "snapshot":
            await ws.send(json.dumps({"op": "snapshot", "objects": await world.snapshot()}))
        elif op == "spawn":
            pose = Pose(**msg.get("pose", {})) if msg.get("pose") else Pose()
            obj = await world.spawn(
                type=msg.get("type", "entity"),
                pose=pose,
                properties=msg.get("properties"),
                owner=msg.get("owner"),
            )
            await ws.send(json.dumps({"op": "ack", "id": obj.id}))
        elif op == "move":
            obj = await world.move(msg["id"], Pose(**msg["pose"]))
            await ws.send(json.dumps({"op": "ack", "id": obj.id}))
        elif op == "remove":
            await world.remove(msg["id"])
            await ws.send(json.dumps({"op": "ack", "id": msg["id"]}))
        else:
            await ws.send(json.dumps({"op": "error", "error": f"unknown op {op!r}"}))
    except KeyError as e:
        await ws.send(json.dumps({"op": "error", "error": f"unknown object {e}"}))
    except (TypeError, ValueError) as e:
        await ws.send(json.dumps({"op": "error", "error": str(e)}))


async def main(host: str, port: int) -> None:
    world = WorldState()
    shell = SpatialShell(world)
    async with websockets.serve(lambda ws: handle(ws, world, shell), host, port):
        print(f"XROS spatial shell listening on ws://{host}:{port}")
        await asyncio.Future()  # run forever


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=8765)
    args = p.parse_args()
    asyncio.run(main(args.host, args.port))