import argparse
import asyncio
import json
import logging
import re
import numpy as np
import torch
import websockets
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)


def extract_molmo_points(molmo_output: str):
    points = []
    for match in re.finditer(
        r'x\d*="\s*([0-9]+(?:\.[0-9]+)?)"\s+y\d*="\s*([0-9]+(?:\.[0-9]+)?)"',
        molmo_output,
    ):
        try:
            p = np.array([float(match.group(1)), float(match.group(2))], dtype=np.float32)
        except ValueError:
            continue

        if np.max(p) > 100:
            continue

        points.append(p / 100.0)

    return points


class Molmo:
    def __init__(self, model_name="allenai/Molmo-7B-D-0924"):
        self.processor = AutoProcessor.from_pretrained(
            model_name,
            trust_remote_code=True,
            torch_dtype="auto",
            device_map="auto",
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            trust_remote_code=True,
            torch_dtype="auto",
            device_map="auto",
        )

    def infer(self, rgb: np.ndarray, prompt: str) -> str:
        image = Image.fromarray(rgb)
        inputs = self.processor.process(images=[image], text=prompt)
        inputs = {k: v.to(self.model.device).unsqueeze(0) for k, v in inputs.items()}

        with torch.inference_mode():
            output = self.model.generate_from_batch(
                inputs,
                GenerationConfig(max_new_tokens=200, stop_strings="<|endoftext|>"),
                tokenizer=self.processor.tokenizer,
            )

        gen_tokens = output[0, inputs["input_ids"].size(1):]
        return self.processor.tokenizer.decode(gen_tokens, skip_special_tokens=True)

    def infer_point(self, rgb: np.ndarray, prompt: str) -> np.ndarray:
        text = self.infer(rgb, prompt)
        log.info(f"Molmo output: {text}")
        points = extract_molmo_points(text)
        return points[0] if points else np.array([0.5, 0.5], dtype=np.float32)


class MolmoWebSocketServer:
    def __init__(self, host="0.0.0.0", port=8765):
        self.host = host
        self.port = port
        self.molmo = Molmo()

    async def handle_client(self, websocket):
        client_id = id(websocket)
        log.info(f"Client connected: {client_id}")

        try:
            async for msg in websocket:
                try:
                    req = json.loads(msg)
                    action = req.get("action")

                    if action != "infer_point":
                        raise ValueError("Only 'infer_point' action is supported")

                    rgb = np.array(req["rgb"], dtype=np.uint8)

                    if "object_name" in req:
                        prompt = f"Point to the center of the {req['object_name']}."
                        label = req["object_name"].replace(" ", "_")
                    elif "prompt" in req:
                        prompt = req["prompt"]
                        label = "custom_prompt"
                    else:
                        raise ValueError("Provide either 'object_name' or 'prompt'")

                    point = self.molmo.infer_point(rgb, prompt)

                    resp = {
                        "status": "ok",
                        "action": "infer_point",
                        "point": point.tolist(),
                    }

                except Exception as e:
                    log.exception("Request error")
                    resp = {"status": "error", "message": str(e)}

                await websocket.send(json.dumps(resp))

        except websockets.exceptions.ConnectionClosed:
            log.info(f"Client disconnected: {client_id}")

    async def start(self):
        log.info(f"Starting Molmo server on ws://{self.host}:{self.port}")
        async with websockets.serve(
            self.handle_client,
            self.host,
            self.port,
            max_size=50 * 1024 * 1024,
        ):
            await asyncio.Future()


def main():
    parser = argparse.ArgumentParser("Molmo Pointing Server")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8765)
    args = parser.parse_args()

    server = MolmoWebSocketServer(
        host=args.host,
        port=args.port,
    )
    asyncio.run(server.start())


if __name__ == "__main__":
    main()
