import json
import numpy as np
import websockets

class MolmoClient:
    def __init__(self, host="localhost", port=8765):
        self.uri = f"ws://{host}:{port}"
        self.websocket = None

    async def connect(self):
        self.websocket = await websockets.connect(
            self.uri,
            max_size=50 * 1024 * 1024,
        )

    async def disconnect(self):
        if self.websocket:
            await self.websocket.close()

    async def infer_point(self, rgb: np.ndarray, object_name: str = None, prompt: str = None):
        if not self.websocket:
            raise RuntimeError("Not connected to server. Call connect() first.")

        req = {
            "action": "infer_point",
            "rgb": rgb.tolist(),
        }

        if object_name:
            req["object_name"] = object_name
        elif prompt:
            req["prompt"] = prompt
        else:
            raise ValueError("Provide either 'object_name' or 'prompt'")

        await self.websocket.send(json.dumps(req))
        resp = json.loads(await self.websocket.recv())

        if resp["status"] == "error":
            raise RuntimeError(f"Server error: {resp['message']}")

        return np.array(resp["point"], dtype=np.float32)

    async def __aenter__(self):
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.disconnect()
