import gradio as gr
import tempfile
from information_extraction import extract_invoice


def process_image(image):
    if image is None:
        return None
    with tempfile.NamedTemporaryFile(suffix='.jpg',delete=True) as tmp:
        image.save(tmp.name)
        result = extract_invoice(tmp.name)
    return result


with gr.Blocks(title="Invoice Extraction with CrewAI") as demo:
    gr.Markdown(
        """
        # Invoice Information Extraction
        Upload an Invoice and get structured JSON
        """
    )

    with gr.Row():
        ## Left column upload invoice
        with gr.Column(scale=1):
            image_input = gr.Image(
                            type ='pil',
                            label = "Upload Invoice Image",
                            height = None,
                            width = None  
            )
            extract_btn = gr.Button("Extract Invoice", variant="primary")

        with gr.Column(scale=1):
            output=gr.JSON(label='Extracted Invoice JSOn')

    extract_btn.click(
        fn=process_image,
        inputs=image_input,
        outputs=output
    )        
    
if __name__=="__main__":
    demo.launch()