import gradio as gr
import requests
import json
import base64

API_URL = "https://image-modal.nebuia.com/extract_id_data"

def send_image_to_api(image_path, json_prompt):
    # Read the image file and convert to base64
    with open(image_path, "rb") as image_file:
        image_base64 = base64.b64encode(image_file.read()).decode('utf-8')
    
    # Parse the JSON prompt
    try:
        json_prompt_dict = json.loads(json_prompt)
    except json.JSONDecodeError:
        return "Error: Invalid JSON prompt"

    # Create a dictionary with the file data and JSON prompt
    files = {
        "file": ("image.jpg", base64.b64decode(image_base64), "image/jpeg")
    }
    data = {
        "json_prompt": json.dumps(json_prompt_dict)
    }
    
    try:
        # Send POST request to the API
        response = requests.post(API_URL, files=files, data=data)
        
        # Check if the request was successful
        if response.status_code == 200:
            # Parse the JSON response
            result = response.json()
            if result.get("success"):
                return json.dumps(result["data"], indent=2)
            else:
                error_message = f"Error in processing:\n{result.get('error', 'Unknown error')}\n"
                error_message += f"Raw output: {result.get('raw_output', 'No raw output available')}"
                return error_message
        else:
            return f"Error: Received status code {response.status_code}\n{response.text}"
    except requests.RequestException as e:
        return f"Error sending request: {e}"

# Define the Gradio interface
def gradio_interface(image_path, json_prompt):
    if image_path is None:
        return "Please upload an image."
    return send_image_to_api(image_path, json_prompt)

# Create the Gradio interface
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Image(type="filepath", label="Sube un ID"),
        gr.Code(label="JSON Prompt", language="json", lines=10, value='{\n    "nombre": ""\n}')
    ],
    outputs=gr.Textbox(label="API Response", lines=10),
    title="NebuIA ID Estructure Extractor",
    description="Sube un ID, modifica la estructura con los datos que deseas extraer y presiona Submit."
)

# Launch the interface
iface.launch()