app.py

"""
DocuMint AI - Invoice & Document Intelligence API
Deploy on Hugging Face Spaces (Gradio) for free hosting + automatic API endpoint.
"""

import os
import json
import time
import threading
import tempfile
import gradio as gr
from PIL import Image

try:
    import spaces
    gpu_decorator = spaces.GPU(duration=60)
except (ImportError, Exception):
    gpu_decorator = lambda f: f

from extraction import extract_invoice, extract_from_pdf, ExtractionResult

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

HF_TOKEN = os.environ.get("HF_TOKEN", "")
APP_TITLE = "DocuMint AI"
APP_DESCRIPTION = """
## AI-Powered Invoice & Document Extraction

Upload an invoice, receipt, or purchase order and get **structured JSON data** instantly.

**How it works:**
1. Upload an image (PNG, JPG) or PDF of your invoice
2. Our AI extracts vendor details, line items, amounts, dates, and tax info
3. Download the structured JSON or integrate via our API

**Extraction Methods:**
- 🧠 **AI Vision** (Primary): Uses Qwen2.5-VL for high-accuracy document understanding
- 📝 **OCR + Patterns** (Fallback): Tesseract OCR with intelligent pattern matching

---
"""

# ---------------------------------------------------------------------------
# Processing Functions
# ---------------------------------------------------------------------------

@gpu_decorator
def process_document(file, use_ocr_only: bool = False):
    """Process an uploaded document and return extraction results."""
    if file is None:
        return "⚠️ Please upload a document first.", "{}"

    start_time = time.time()

    try:
        file_path = file if isinstance(file, str) else file.name
        file_ext = os.path.splitext(file_path)[1].lower()

        if file_ext == ".pdf":
            with open(file_path, "rb") as f:
                pdf_bytes = f.read()
            results = extract_from_pdf(
                pdf_bytes,
                hf_token=HF_TOKEN,
                force_ocr=use_ocr_only,
            )
            if not results:
                return "⚠️ Could not extract data from PDF.", "{}"

            # Merge multi-page results or return first page
            result = results[0]
            if len(results) > 1:
                # For multi-page invoices, merge line items
                for r in results[1:]:
                    result.line_items.extend(r.line_items)

        else:
            image = Image.open(file_path)
            result = extract_invoice(
                image,
                hf_token=HF_TOKEN,
                force_ocr=use_ocr_only,
            )

        total_time = int((time.time() - start_time) * 1000)
        result.processing_time_ms = total_time

        # Check for informative failure modes
        if result.extraction_method.startswith("failed"):
            method = result.extraction_method
            if "no_token" in method:
                warning = (
                    "⚠️ **Extraction failed:** OCR could not read this document "
                    "(likely handwritten or low quality).\n\n"
                    "**AI Vision is not enabled.** Set the `HF_TOKEN` Space Secret "
                    "to unlock Qwen2.5-VL, which handles handwritten text, "
                    "non-English documents, and complex layouts with much higher accuracy."
                )
            elif "vlm" in method:
                warning = (
                    "⚠️ **Extraction failed:** AI Vision models could not process this document, "
                    "and OCR could not extract readable text either.\n\n"
                    f"**Debug info:** `{method}`\n\n"
                    "This can happen when HF Inference API models are loading or rate-limited. "
                    "Try again in a minute, or check that your `HF_TOKEN` has valid API access."
                )
            else:
                warning = (
                    "⚠️ **Extraction failed:** Could not extract readable text from this document. "
                    "The image may be too blurry, rotated, or handwritten for OCR.\n\n"
                    "Try uploading a clearer scan, or uncheck 'OCR only' to use AI Vision."
                )
            json_output = json.dumps(result.to_dict(), indent=2, default=str)
            return warning, json_output

        # Format summary
        summary = format_summary(result)
        json_output = json.dumps(result.to_dict(), indent=2, default=str)

        return summary, json_output

    except Exception as e:
        return f"❌ Error: {str(e)}", "{}"


def format_summary(result: ExtractionResult) -> str:
    """Format extraction result as a readable summary."""
    lines = []
    lines.append(f"### Extraction Results")
    lines.append(f"**Method:** {result.extraction_method}")
    lines.append(f"**Confidence:** {result.confidence:.0%}")
    lines.append(f"**Processing Time:** {result.processing_time_ms}ms")
    lines.append("")

    if result.vendor_name:
        lines.append(f"**Vendor:** {result.vendor_name}")
    if result.invoice_number:
        lines.append(f"**Invoice #:** {result.invoice_number}")
    if result.invoice_date:
        lines.append(f"**Date:** {result.invoice_date}")
    if result.due_date:
        lines.append(f"**Due Date:** {result.due_date}")
    if result.currency:
        lines.append(f"**Currency:** {result.currency}")

    lines.append("")

    if result.line_items:
        lines.append("**Line Items:**")
        lines.append("")
        lines.append("| Description | Qty | Unit Price | Amount |")
        lines.append("|---|---|---|---|")
        for item in result.line_items:
            desc = item.description if isinstance(item, object) and hasattr(item, 'description') else item.get('description', '')
            qty = item.quantity if hasattr(item, 'quantity') else item.get('quantity', '')
            price = item.unit_price if hasattr(item, 'unit_price') else item.get('unit_price', '')
            amt = item.amount if hasattr(item, 'amount') else item.get('amount', '')
            lines.append(f"| {desc} | {qty or '-'} | {price or '-'} | {amt or '-'} |")
        lines.append("")

    if result.subtotal is not None:
        lines.append(f"**Subtotal:** {result.currency} {result.subtotal:,.2f}")
    if result.tax_amount is not None:
        tax_str = f"**Tax:** {result.currency} {result.tax_amount:,.2f}"
        if result.tax_rate:
            tax_str += f" ({result.tax_rate})"
        lines.append(tax_str)
    if result.total_amount is not None:
        lines.append(f"**Total:** {result.currency} {result.total_amount:,.2f}")

    if result.customer_name:
        lines.append(f"\n**Billed To:** {result.customer_name}")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Gradio Interface
# ---------------------------------------------------------------------------

def create_demo():
    """Create the Gradio demo interface."""

    with gr.Blocks(
        title=APP_TITLE,
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="slate",
        ),
        css="""
        .main-header { text-align: center; margin-bottom: 1rem; }
        .main-header h1 { font-size: 2rem; color: #1e40af; }
        .output-json { font-family: monospace; font-size: 0.85rem; }
        footer { display: none !important; }
        """
    ) as demo:

        gr.HTML("""
        <div class="main-header">
            <h1>📄 DocuMint AI</h1>
            <p style="color: #64748b; font-size: 1.1rem;">
                Extract structured data from invoices, receipts & purchase orders
            </p>
        </div>
        """)

        gr.Markdown(APP_DESCRIPTION)

        with gr.Row():
            with gr.Column(scale=1):
                file_input = gr.File(
                    label="Upload Invoice (Image or PDF)",
                    file_types=["image", ".pdf"],
                    type="filepath",
                )
                ocr_only = gr.Checkbox(
                    label="Use OCR only (skip AI vision model)",
                    value=False,
                    info="Check this if the AI model is unavailable or you want faster processing",
                )
                extract_btn = gr.Button(
                    "🔍 Extract Data",
                    variant="primary",
                    size="lg",
                )

            with gr.Column(scale=1):
                summary_output = gr.Markdown(
                    label="Extraction Summary",
                    value="*Upload a document and click Extract to begin.*",
                )

        with gr.Row():
            json_output = gr.Code(
                label="Structured JSON Output",
                language="json",
                lines=20,
            )

        # API usage instructions
        with gr.Accordion("🔌 API Usage", open=False):
            gr.Markdown("""
### Use DocuMint AI as an API

This Gradio app is automatically exposed as an API endpoint on Hugging Face Spaces.

**Python Example:**
```python
from gradio_client import Client

client = Client("YOUR_SPACE_URL")
result = client.predict(
    file="path/to/invoice.png",
    use_ocr_only=False,
    api_name="/predict"
)
summary, json_data = result
print(json_data)
```

**cURL Example:**
```bash
curl -X POST "YOUR_SPACE_URL/api/predict" \\
  -H "Content-Type: multipart/form-data" \\
  -F "file=@invoice.png" \\
  -F "use_ocr_only=false"
```

**Response Format:**
```json
{
  "vendor_name": "Acme Corp",
  "invoice_number": "INV-2024-001",
  "invoice_date": "2024-01-15",
  "total_amount": 1250.00,
  "currency": "USD",
  "line_items": [...],
  "tax_amount": 225.00,
  "confidence": 0.92
}
```
            """)

        # Pricing section
        with gr.Accordion("💰 Pricing Plans", open=False):
            gr.Markdown("""
| Plan | Pages/Month | Price | Per Page |
|------|------------|-------|----------|
| **Free** | 50 | $0 | $0.00 |
| **Starter** | 500 | $29/mo | $0.058 |
| **Pro** | 2,000 | $79/mo | $0.040 |
| **Business** | 10,000 | $199/mo | $0.020 |
| **Enterprise** | Unlimited | Custom | Contact us |

*All plans include JSON + CSV export, API access, and email support.*
*Pro and above include priority processing and webhook notifications.*

📧 Contact: **documint.ai@gmail.com** for enterprise pricing and custom integrations.
            """)

        extract_btn.click(
            fn=process_document,
            inputs=[file_input, ocr_only],
            outputs=[summary_output, json_output],
        )

    return demo


# ---------------------------------------------------------------------------
# Entry Point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    demo = create_demo()

    # Launch Gradio with prevent_thread_lock so we can add custom API routes
    # to Gradio's internal FastAPI app. This preserves ZeroGPU compatibility
    # (demo.launch handles @spaces.GPU registration) while also serving our
    # REST endpoints for the React frontend.
    app, local_url, _ = demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        ssr_mode=False,
        prevent_thread_lock=True,
    )

    # Inject our REST API routes into Gradio's FastAPI app.
    # Gradio's app already includes CORSMiddleware (allow_origins=["*"]),
    # so our routes inherit CORS handling — no extra middleware needed.
    from api import api_router
    from api_runs import runs_router
    from api_prompts import prompts_router
    app.include_router(api_router)
    app.include_router(runs_router)
    app.include_router(prompts_router)

    print(f"✅ REST API routes mounted at {local_url}api/v1/")

    # Keep the main thread alive (uvicorn runs in a background thread)
    threading.Event().wait()