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

def process(json_input):
    try:
        data = json.loads(json_input)

        # Decode Base64
        img_bytes = base64.b64decode(data["image_b64"])
        img = Image.open(io.BytesIO(img_bytes))

        # You can do CV/LLM here later
        response = {
            "received": True,
            "robot_id": data.get("robot_id", "unknown"),
            "image_size": img.size
        }

        return json.dumps(response)

    except Exception as e:
        return json.dumps({"error": str(e)})

# Gradio interface
demo = gr.Interface(
    fn=process,
    inputs=gr.Textbox(lines=10, label="Input JSON from Jetson"),
    outputs=gr.Textbox(lines=10, label="Reply to Jetson"),
)

demo.launch()