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

# Configure Gemini API 
API_KEY = os.getenv("GOOGLE_API_KEY")

if API_KEY:
    genai.configure(api_key=API_KEY)

# Function: analyze invoice with Gemini
def analyze_invoice(user_query, uploaded_image):
    if not API_KEY:
        return "❌ Error: GOOGLE_API_KEY not configured. Please add your API key in Settings > Repository secrets."
    
    if uploaded_image is None:
        return "⚠️ Please upload an invoice image."
    
    if not user_query.strip():
        return "⚠️ Please enter a question about the invoice."
    
    try:
        # Convert uploaded image to bytes
        with open(uploaded_image, "rb") as f:
            image_bytes = f.read()
        
        image_data = {
            "mime_type": "image/jpeg",
            "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
        
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Gradio UI
with gr.Blocks(title="📄 Invoice Reader - Gemini AI") as demo:
    gr.Markdown("# 📄 Invoice Reader using Gemini AI")
    gr.Markdown("### Upload an invoice and ask questions about it!")
    
    if not API_KEY:
        gr.Markdown("❌ **Status**: Please add GOOGLE_API_KEY in Settings > Repository secrets")
    
    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?",
                lines=2
            )
            analyze_btn = gr.Button("🔍 Analyze Invoice", variant="primary")
        
        with gr.Column(scale=2):
            output = gr.Textbox(label="Gemini Response", lines=15)
    
    # Examples
    gr.Examples(
        examples=[
            ["What is the total amount?"],
            ["Who is the vendor?"],
            ["Extract all line items"],
            ["What is the invoice number?"],
        ],
        inputs=user_query
    )
    
    analyze_btn.click(
        fn=analyze_invoice,
        inputs=[user_query, invoice_image],
        outputs=output
    )

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