import gradio as gr
import base64
import httpx
import os

API_URL = "http://127.0.0.1:7860/api/document-analyze"
API_KEY = os.getenv("API_KEY", "sk_track2_987654321")

async def analyze(file):
    if file is None:
        return "Please upload a document."
    
    file_name = os.path.basename(file.name)
    ext = file_name.split('.')[-1].lower()
    file_type = "image" if ext in ['jpg', 'jpeg', 'png'] else ext
        
    with open(file.name, "rb") as f:
        encoded_string = base64.b64encode(f.read()).decode('utf-8')

    payload = {
        "fileName": file_name,
        "fileType": file_type,
        "fileBase64": encoded_string
    }
        
    async with httpx.AsyncClient(timeout=60.0) as client:
        try:
            response = await client.post(API_URL, json=payload, headers={"x-api-key": API_KEY})
            res = response.json()
            
            # Note: We use .get() to handle cases where 'status' might be missing
            report = f"## ✅ Analysis Complete: {res.get('fileName', file_name)}\n\n"
            report += f"### 🤖 AI Summary\n{res.get('summary', 'No summary')}\n\n"
            report += f"### 📊 Sentiment: **{res.get('sentiment', 'N/A')}**\n\n"
            
            entities = res.get('entities', {})
            report += "### 🔍 Key Entities Found\n"
            report += f"- **Names:** {', '.join(entities.get('names', [])) or 'None'}\n"
            report += f"- **Organizations:** {', '.join(entities.get('organizations', [])) or 'None'}\n"
            report += f"- **Dates:** {', '.join(entities.get('dates', [])) or 'None'}\n"
            report += f"- **Financials:** {', '.join(entities.get('amounts', [])) or 'None'}\n"
            return report
        except Exception as e:
            return f"### ❌ Connection Error\n{str(e)}"

# Define the interface as 'demo'
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue")) as demo:
    gr.Markdown("# 🏢 Intelligent Document Processing (IDP) System")
    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(label="Upload Document", file_types=[".pdf", ".docx", ".jpg", ".png"])
            submit_btn = gr.Button("Analyze Now", variant="primary")
        with gr.Column(scale=2):
            output_display = gr.Markdown("Results will appear here...")
    
    submit_btn.click(fn=analyze, inputs=file_input, outputs=output_display)