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

# The input will now be a raw text string representing the JSON data
def process(json_text_input): 
    try:
        # Manually parse the incoming text string into a Python dictionary
        data = json.loads(json_text_input) 

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

        response = {
            "received": True,
            "robot_id": data.get("robot_id", "unknown"),
            "image_size": img.size
        }
        return response

    except Exception as e:
        # Return error as a JSON dictionary which Gradio handles
        return {"error": str(e)}

demo = gr.Interface(
    fn=process,
    # Change input type to Textbox
    inputs=gr.Textbox(label="Input JSON as Plain Text"),
    outputs=gr.JSON(label="Reply to Jetson"),
)

demo.launch()