import gradio as gr
from PIL import Image
import google.generativeai as genai
import os

# Configure Gemini API (from Hugging Face Secrets)
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# Function: analyze invoice with Gemini
def analyze_invoice(user_query, uploaded_image):
    if uploaded_image is None:
        return "⚠️ Please upload an invoice image."
    if not user_query.strip():
        return "⚠️ Please enter a question about the invoice."
    
    # Convert uploaded image to bytes
    with open(uploaded_image, "rb") as f:
        image_bytes = f.read()

    image_data = {
        "mime_type": "image/jpeg",  # Gradio ensures jpg/png
        "data": image_bytes
    }

    # Prompt for Gemini
    prompt = f"""
    You are an expert in understanding invoices.
    Extract key fields (Invoice Number, Date, Vendor, Customer, 
    Line Items with Description, Quantity, Unit Price, Total, and Taxes).
    Then answer this query: {user_query}.
    """

    model = genai.GenerativeModel("gemini-1.5-flash")
    response = model.generate_content([prompt, image_data])
    return response.text

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 📄 Invoice Reader using Gemini API (Gradio)")

    with gr.Row():
        with gr.Column(scale=1):
            invoice_image = gr.Image(type="filepath", label="Upload Invoice (JPG/PNG)")
            user_query = gr.Textbox(label="Ask about the invoice", placeholder="What is the total amount?")
            analyze_btn = gr.Button("Analyze Invoice")
        
        with gr.Column(scale=2):
            output = gr.Textbox(label="Gemini Response", lines=12)

    analyze_btn.click(
        fn=analyze_invoice,
        inputs=[user_query, invoice_image],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch()