import gradio as gr
import requests
import json
import base64
import os


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

def send_file_to_api(file_path, json_prompt):
    # Read the file and convert to base64
    with open(file_path, "rb") as file:
        file_base64 = base64.b64encode(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"
    
    # Determine file type
    file_extension = os.path.splitext(file_path)[1].lower()
    if file_extension in ['.jpg', '.jpeg', '.png']:
        mime_type = "image/jpeg"
    elif file_extension == '.pdf':
        mime_type = "application/pdf"
    else:
        return "Error: Unsupported file type"
    
    # Create a dictionary with the file data and JSON prompt
    files = {
        "file": (os.path.basename(file_path), base64.b64decode(file_base64), mime_type)
    }
    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(file, json_prompt):
    if file is None:
        return "Please upload an image or PDF file."
    return send_file_to_api(file.name, json_prompt)

# Custom color for the theme
custom_purple = "#7f56d9"

# Create the Gradio interface using Blocks
with gr.Blocks(theme=gr.themes.Default(primary_hue="purple", secondary_hue="purple")) as demo:
    gr.Markdown(
        """
        <div style="text-align: center;">
            <img src="https://copilot.nebuia.com/images/logo_white.png" style="max-height: 400px; max-width: 120px; object-fit: contain;">
        </div>

        ## NebuIA ID Structure Extractor
        Sube un ID (imagen o PDF), modifica la estructura con los datos que deseas extraer y presiona Submit.
        """
    )
    with gr.Row():
        file_input = gr.File(label="Sube un ID (imagen o PDF)", type="filepath")
        json_input = gr.Code(label="JSON Prompt", language="json", lines=10, value='{\n    "nombre": ""\n}')
    
    output = gr.Textbox(label="API Response", lines=10)
    submit_btn = gr.Button("Submit")
    
    submit_btn.click(fn=gradio_interface, inputs=[file_input, json_input], outputs=output)

# Launch the interface
demo.launch()