import gradio as gr
import base64
from PIL import Image
import io
import json
import requests

HF_VLM_API = "https://api-inference.huggingface.co/models/Qwen/Qwen2-VL-7B-Instruct"
HF_TOKEN = "HF_CV_ROBOT_TOKEN"  # HF Token

def call_vlm_api(img: Image):
    # encode image to bytes
    buf = io.BytesIO()
    img.save(buf, format="JPEG")
    img_bytes = buf.getvalue()

    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    payload = {"inputs": [{"image": img_bytes, "text": "Describe the image in detail."}]}
    
    resp = requests.post(HF_VLM_API, headers=headers, json=payload, timeout=60)
    if resp.status_code == 200:
        return resp.json()[0].get("generated_text", "")
    else:
        return f"VLM API error: {resp.status_code}"

def process(payload: dict):
    try:
        img_bytes = base64.b64decode(payload["image_b64"])
        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")

        vlm_text = call_vlm_api(img)

        reply = {
            "received": True,
            "robot_id": payload.get("robot_id", "unknown"),
            "size": img.size,
            "vllm_analysis": vlm_text
        }
        return reply
    except Exception as e:
        return {"error": str(e)}

demo = gr.Interface(
    fn=process,
    inputs=gr.JSON(label="Input Payload (Dict format)"),
    outputs=gr.JSON(label="Reply to Jetson"),
    api_name="predict"
)

if __name__ == "__main__":
    demo.launch(mcp_server=True)