import gradio as gr
import os
from dotenv import load_dotenv
import pandas as pd
from groq import Groq
from PIL import Image
import base64
import io
import openpyxl
from datetime import datetime
import httpx

# Load environment variables from .env file
load_dotenv()

def create_groq_client():
    api_key = os.environ.get("GROQ_API_KEY", "")
    return httpx.Client(
        base_url="https://api.groq.com/openai/v1",
        headers={"Authorization": f"Bearer {api_key}"}
    )

def encode_image_to_base64(image_path):
    """Convert image to base64 string"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def extract_invoice_details(image):
    """Extract invoice details using Groq's vision model"""
    # Save the uploaded image temporarily
    temp_path = "temp_invoice.png"
    image.save(temp_path)
    
    # Convert image to base64
    base64_image = encode_image_to_base64(temp_path)
    
    # Remove temporary file
    os.remove(temp_path)
    
    # Prepare the prompt
    prompt = """Analyze this invoice image and provide ONLY ONE dictionary with the following format, including all line items. Remove any special characters (*, #, $) and format numbers as plain decimal values:

    {
        "Invoice Number": "inv-00",  # Remove special chars, keep alphanumeric only
        "Invoice Date": "07/07/2025",  # Use MM/DD/YYYY format
        "Items": [
            {
                "Item Name": "Product 1",  # Clean text only
                "Price/Rate": "40.00",  # Numeric only
                "Quantity": "2",  # Numeric only
                "Amount": "80.00"  # Numeric only
            }
        ],
        "Total Invoice Value": "2555.00"  # Numeric only, total amount
    }

    Provide ONLY the dictionary, no additional text or formatting."""
    
    # Make API call to Groq
    client = create_groq_client()
    response = client.post(
        "/chat/completions",
        json={
            "model": "llama-3.2-90b-vision-preview",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{base64_image}"
                            }
                        },
                        {
                            "type": "text",
                            "text": prompt
                        }
                    ]
                }
            ]
        }
    )
    response_data = response.json()
    return response_data['choices'][0]['message']['content']

def parse_response(response_text):
    """Parse the model's response into structured data"""
    try:
        # Find the dictionary part of the response
        start_idx = response_text.find('{')
        end_idx = response_text.rfind('}') + 1
        if start_idx != -1 and end_idx != -1:
            dict_str = response_text[start_idx:end_idx]
            # Safely evaluate the dictionary string
            data = eval(dict_str)
            
            # Create rows for each item in the items list
            rows = []
            for item in data.get('Items', []):
                row = {
                    'Invoice Number': data.get('Invoice Number', ''),
                    'Invoice Date': data.get('Invoice Date', ''),
                    'Item Name': item.get('Item Name', ''),
                    'Price/Rate': item.get('Price/Rate', ''),
                    'Quantity': item.get('Quantity', ''),
                    'Amount': item.get('Amount', ''),
                    'Total Invoice Value': data.get('Total Invoice Value', '')
                }
                rows.append(row)
            
            return rows
    except Exception as e:
        print(f"Error parsing response: {e}")
        return [{
            'Invoice Number': '',
            'Invoice Date': '',
            'Item Name': '',
            'Price/Rate': '',
            'Quantity': '',
            'Amount': '',
            'Total Invoice Value': ''
        }]

def save_to_excel(data_rows):
    """Save cleaned extracted data to Excel file"""
    excel_file = "invoice_data.xlsx"
    
    # Create new DataFrame with the current data only
    df = pd.DataFrame(data_rows, columns=[
        'Invoice Number', 'Invoice Date', 'Item Name',
        'Price/Rate', 'Quantity', 'Amount', 'Total Invoice Value'
    ])
    
    # Apply number formatting for currency columns
    currency_columns = ['Price/Rate', 'Amount', 'Total Invoice Value']
    for col in currency_columns:
        df[col] = pd.to_numeric(df[col], errors='ignore')
    
    # Save to Excel with formatting
    with pd.ExcelWriter(excel_file, engine='openpyxl') as writer:
        df.to_excel(writer, index=False, sheet_name='Invoice Data')
        
        # Get the workbook and worksheet
        workbook = writer.book
        worksheet = writer.sheets['Invoice Data']
        
        # Apply currency formatting to relevant columns
        for col_idx, col_name in enumerate(df.columns):
            if col_name in currency_columns:
                for row in range(2, len(df) + 2):  # Start from row 2 to skip header
                    cell = worksheet.cell(row=row, column=col_idx + 1)
                    cell.number_format = '$#,##0.00'
    
    return excel_file

def process_invoice(image):
    """Main function to process invoice image"""
    try:
        # Extract text from image
        extracted_text = extract_invoice_details(image)
        
        # Parse the response
        data = parse_response(extracted_text)
        
        # Save to Excel
        excel_path = save_to_excel(data)
        
        return (
            f"Successfully processed invoice!\n\n"
            f"Extracted Information:\n{extracted_text}",
            excel_path
        )
    
    except Exception as e:
        return f"Error processing invoice: {str(e)}", None

# Create Gradio interface
iface = gr.Interface(
    fn=process_invoice,
    inputs=gr.Image(type="pil", label="Upload Handwritten Invoice"),
    outputs=[
        gr.Textbox(label="Processing Result"),
        gr.File(label="Download Excel File")
    ],
    title="Handwritten Invoice Processor",
    description="Upload a handwritten invoice to extract key information and save it to Excel.",
    examples=[],
    theme=gr.themes.Base()
)

# Launch the application
if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)