app.py

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
from pathlib import Path
from typing import List
import requests
import base64
import json
import re
import fitz  # pymupdf — no poppler required

app = FastAPI(
    title="Invoice OCR API",
    description="Two-step pipeline: nemotron-ocr-v1 → nvidia-nemotron-nano-9b-v2 for Tax Invoice extraction. Supports images AND multi-page PDFs.",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Configuration ─────────────────────────────────────────────────────────────
NVIDIA_API_KEY = "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"

OCR_URL = "https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-ocr-v1"
LLM_URL = "https://integrate.api.nvidia.com/v1/chat/completions"
LLM_MODEL = "nvidia/nvidia-nemotron-nano-9b-v2"

OCR_HEADERS = {"Authorization": f"Bearer {NVIDIA_API_KEY}", "Accept": "application/json"}
LLM_HEADERS = {"Authorization": f"Bearer {NVIDIA_API_KEY}", "Content-Type": "application/json"}

PDF_MAX_PAGES = 10

# ── System prompt ─────────────────────────────────────────────────────────────
INVOICE_SYSTEM_PROMPT = """You are a Tax Invoice data extraction assistant for Indian GST invoices.
You will receive OCR text from a tax invoice image. Return ONLY a valid JSON object. No markdown, no explanation.

JSON schema (return exactly this):
{
  "invoice_number": "invoice number e.g. ACMPL/01/19-20 (string)",
  "eway_bill_number": "e-Way Bill No if present (string)",
  "invoice_date": "date e.g. 18-Apr-2019 (string)",
  "mode_of_payment": "cash/credit/UPI/bank transfer etc (string)",
  "supplier_ref": "supplier reference number (string)",
  "buyer_order_number": "buyer's PO number (string)",
  "dispatch_document_number": "dispatch doc number (string)",
  "dispatched_through": "courier or transport name (string)",
  "destination": "delivery destination (string)",
  "delivery_note": "delivery note number (string)",

  "vendor_name": "name of selling company e.g. Ace Mobile Manufacturer Pvt Ltd (string)",
  "vendor_address": "full address of vendor (string)",
  "vendor_gstin": "15-char GSTIN of vendor e.g. 09AABCS1429B1ZS (string)",
  "vendor_state": "state name and code e.g. Uttar Pradesh, Code: 09 (string)",
  "vendor_email": "email of vendor (string)",

  "buyer_name": "name of buyer/customer e.g. The Mobile Planet (string)",
  "buyer_address": "full address of buyer (string)",
  "buyer_gstin": "15-char GSTIN of buyer e.g. 09AAGCA1654H1ZQ (string)",
  "buyer_state": "state name and code of buyer (string)",

  "items": [
    {
      "sl_no": "serial number (string)",
      "description": "description of goods (string)",
      "batch": "batch number if present (string)",
      "hsn_sac": "HSN or SAC code (string)",
      "quantity": "quantity with unit e.g. 500 Nos (string)",
      "rate": "rate per unit e.g. 6000.00 (string)",
      "per": "unit type e.g. Nos (string)",
      "amount": "line total e.g. 30,00,000.00 (string)"
    }
  ],

  "taxable_value": "total taxable amount before tax (string)",
  "cgst_rate": "CGST rate percentage e.g. 6% (string)",
  "cgst_amount": "CGST amount (string)",
  "sgst_rate": "SGST rate percentage e.g. 6% (string)",
  "sgst_amount": "SGST amount (string)",
  "igst_rate": "IGST rate if applicable (string)",
  "igst_amount": "IGST amount if applicable (string)",
  "output_cgst": "Output CGST amount (string)",
  "output_sgst": "Output SGST amount (string)",
  "total_tax_amount": "total tax amount (string)",
  "grand_total": "final invoice total e.g. 96,32,000.00 (string)",
  "amount_in_words": "amount in words e.g. INR Ninety Six Lakh Thirty Two Thousand Only (string)",
  "tax_amount_in_words": "tax amount in words (string)",

  "hsn_summary": [
    {
      "hsn_sac": "HSN code (string)",
      "taxable_value": "taxable value for this HSN (string)",
      "cgst_rate": "CGST rate (string)",
      "cgst_amount": "CGST amount (string)",
      "sgst_rate": "SGST rate (string)",
      "sgst_amount": "SGST amount (string)",
      "total_tax": "total tax for this HSN (string)"
    }
  ],

  "declaration": "declaration text at bottom (string)",
  "authorised_signatory": "authorised signatory label (string)",
  "is_computer_generated": true
}

CRITICAL RULES:
- invoice_number: look for Invoice No., Bill No., Ref No. near the top right area
- vendor_name: the company at the TOP of the invoice, usually with logo
- buyer_name: look for 'Buyer', 'Bill To', 'Sold To' section
- GSTIN: exactly 15 characters, mix of letters and digits e.g. 09AABCS1429B1ZS
- items: extract EVERY line item row in the goods table including batch info
- amounts: keep exact format with commas e.g. 30,00,000.00
- hsn_summary: extract the tax summary table at the bottom (HSN/SAC wise breakdown)
- output_cgst / output_sgst: look for 'Output CGST' and 'Output SGST' labels in totals
- grand_total: the final TOTAL amount, look for ₹ symbol
- amount_in_words: the spelled-out amount e.g. 'INR Ninety Six Lakh...'
- If a field is not found, use "" for strings, [] for arrays, false for booleans"""


# ── PDF → PNG list (PyMuPDF — no poppler required) ────────────────────────────

def pdf_bytes_to_png_list(pdf_bytes: bytes, dpi: int = 200) -> list[bytes]:
    """Convert every page of a PDF into PNG bytes using PyMuPDF."""
    try:
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    except Exception as exc:
        raise HTTPException(422, f"Could not read PDF: {exc}")

    if doc.page_count == 0:
        raise HTTPException(422, "PDF has no pages or could not be rendered.")

    if doc.page_count > PDF_MAX_PAGES:
        raise HTTPException(
            400,
            f"PDF has {doc.page_count} pages; maximum allowed is {PDF_MAX_PAGES}. "
            "Please send a trimmed PDF.",
        )

    result = []
    for page in doc:
        pix = page.get_pixmap(dpi=dpi)
        result.append(pix.tobytes("png"))
    doc.close()
    return result


# ── OCR helpers ───────────────────────────────────────────────────────────────

def extract_all_text_sorted(ocr_json: dict) -> tuple[str, list]:
    """Sort OCR detections spatially top-to-bottom, left-to-right in bands."""
    data = ocr_json.get("data", [])
    detections = data[0].get("text_detections", []) if data else ocr_json.get("text_detections", [])

    items = []
    for det in detections:
        if not isinstance(det, dict):
            continue
        if "text_prediction" in det:
            text = det["text_prediction"].get("text", "").strip()
        else:
            text = det.get("text", "").strip()
        if not text:
            continue
        pts = det.get("bounding_box", {}).get("points", [])
        y = sum(p["y"] for p in pts) / len(pts) if pts else 0
        x = min(p["x"] for p in pts) if pts else 0
        items.append({"text": text, "y": y, "x": x})

    BAND = 0.012
    items.sort(key=lambda d: (round(d["y"] / BAND), d["x"]))
    full_text = "\n".join(i["text"] for i in items)
    return full_text, items


def run_ocr_on_bytes(image_bytes: bytes, page_label: str = "") -> tuple[str, list]:
    """Run OCR on raw image bytes. Returns (text, detections)."""
    image_b64 = base64.b64encode(image_bytes).decode()

    if len(image_b64) >= 1_000_000:
        raise HTTPException(
            status_code=413,
            detail=f"Image too large{' (page ' + page_label + ')' if page_label else ''}. Resize and retry."
        )

    payload = {"input": [{"type": "image_url", "url": f"data:image/png;base64,{image_b64}"}]}
    try:
        resp = requests.post(OCR_URL, headers=OCR_HEADERS, json=payload, timeout=30)
        resp.raise_for_status()
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=502, detail=f"NVIDIA OCR error: {str(e)}")

    ocr_json = resp.json()
    text, items = extract_all_text_sorted(ocr_json)
    label = f"page {page_label} " if page_label else ""
    print(f"OCR {label}({len(text)} chars):\n{text[:400]}\n{'='*60}")
    return text, items


async def run_ocr(file: UploadFile) -> tuple[str, list]:
    """
    Read the uploaded file.
    - If PDF  → convert each page to PNG via PyMuPDF, OCR all pages, concatenate.
    - If image → OCR directly.
    Returns (combined_text, detections_of_first_page).
    """
    content = await file.read()
    content_type = (file.content_type or "").lower()
    filename = (file.filename or "").lower()

    is_pdf = (
        content[:4] == b"%PDF"
        or content_type == "application/pdf"
        or filename.endswith(".pdf")
    )

    if is_pdf:
        print(f"PDF detected ({len(content)} bytes). Converting pages to images…")
        page_images = pdf_bytes_to_png_list(content)

        all_texts: list[str] = []
        first_detections: list = []

        for i, img_bytes in enumerate(page_images, start=1):
            page_text, detections = run_ocr_on_bytes(img_bytes, page_label=str(i))
            if page_text.strip():
                all_texts.append(f"--- Page {i} ---\n{page_text}")
            if i == 1:
                first_detections = detections

        combined = "\n\n".join(all_texts)
        print(f"Total combined OCR text: {len(combined)} chars across {len(page_images)} page(s)")
        return combined, first_detections

    # Regular image path
    return run_ocr_on_bytes(content)


# ── LLM ───────────────────────────────────────────────────────────────────────

def call_llm(ocr_text: str) -> dict:
    payload = {
        "model": LLM_MODEL,
        "max_tokens": 5000,
        "temperature": 0.1,
        "top_p": 0.9,
        "messages": [
            {"role": "system", "content": INVOICE_SYSTEM_PROMPT},
            {
                "role": "user",
                "content": (
                    f"OCR text from tax invoice:\n\n{ocr_text}\n\n"
                    "Return ONLY the complete JSON object."
                ),
            },
        ],
    }

    try:
        resp = requests.post(LLM_URL, headers=LLM_HEADERS, json=payload, timeout=200)
        resp.raise_for_status()
        llm_json = resp.json()
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=502, detail=f"NVIDIA LLM error: {str(e)}")

    choice = llm_json.get("choices", [{}])[0]
    raw = choice.get("message", {}).get("content", "")
    finish = choice.get("finish_reason", "")
    print(f"LLM finish={finish}\nRaw (first 600):\n{raw[:600]}\n{'='*60}")

    if not raw:
        raise HTTPException(status_code=502, detail="LLM returned empty response")

    cleaned = re.sub(r"```json\s*", "", raw, flags=re.IGNORECASE)
    cleaned = re.sub(r"```\s*", "", cleaned).strip()

    try:
        parsed = json.loads(cleaned)
        if isinstance(parsed, dict):
            return parsed
    except json.JSONDecodeError:
        pass

    match = re.search(r"\{[\s\S]*\}", cleaned)
    if match:
        try:
            parsed = json.loads(match.group(0))
            if isinstance(parsed, dict):
                return parsed
        except json.JSONDecodeError:
            pass

    patched = cleaned.rstrip().rstrip(",")
    open_braces   = patched.count("{") - patched.count("}")
    open_brackets = patched.count("[") - patched.count("]")
    patched += "]" * max(0, open_brackets) + "}" * max(0, open_braces)
    try:
        parsed = json.loads(patched)
        if isinstance(parsed, dict):
            print("WARNING: used bracket-patching to fix truncated JSON")
            return parsed
    except json.JSONDecodeError:
        pass

    raise HTTPException(
        status_code=502,
        detail=f"JSON parse failed (finish={finish}). Preview: {raw[:400]}"
    )


# ── Pydantic models ───────────────────────────────────────────────────────────

class LineItem(BaseModel):
    sl_no: str
    description: str
    batch: str
    hsn_sac: str
    quantity: str
    rate: str
    per: str
    amount: str

class HSNSummary(BaseModel):
    hsn_sac: str
    taxable_value: str
    cgst_rate: str
    cgst_amount: str
    sgst_rate: str
    sgst_amount: str
    total_tax: str

class InvoiceData(BaseModel):
    invoice_number: str
    eway_bill_number: str
    invoice_date: str
    mode_of_payment: str
    supplier_ref: str
    buyer_order_number: str
    dispatch_document_number: str
    dispatched_through: str
    destination: str
    delivery_note: str

    vendor_name: str
    vendor_address: str
    vendor_gstin: str
    vendor_state: str
    vendor_email: str

    buyer_name: str
    buyer_address: str
    buyer_gstin: str
    buyer_state: str

    items: List[LineItem]

    taxable_value: str
    cgst_rate: str
    cgst_amount: str
    sgst_rate: str
    sgst_amount: str
    igst_rate: str
    igst_amount: str
    output_cgst: str
    output_sgst: str
    total_tax_amount: str
    grand_total: str
    amount_in_words: str
    tax_amount_in_words: str

    hsn_summary: List[HSNSummary]

    declaration: str
    authorised_signatory: str
    is_computer_generated: bool
    source_pages: int = 1


# ── Endpoints ──────────────────────────────────────────────────────────────────

ALLOWED_TYPES = {
    "image/jpeg", "image/jpg", "image/png", "image/webp", "image/gif",
    "application/pdf", "application/x-pdf",
}

@app.post("/extract-invoice", response_model=InvoiceData)
async def extract_invoice(file: UploadFile = File(...)):
    """
    Upload a tax invoice image (JPEG/PNG/WEBP) or PDF →
    structured JSON with all fields.
    """
    content_type = (file.content_type or "").lower()
    filename = (file.filename or "").lower()
    is_pdf = content_type in ("application/pdf", "application/x-pdf") or filename.endswith(".pdf")

    if content_type and content_type not in ALLOWED_TYPES and not is_pdf:
        raise HTTPException(status_code=415, detail=f"Unsupported type: {file.content_type}. Accepted: JPEG, PNG, WebP, PDF.")

    ocr_text, _ = await run_ocr(file)
    if not ocr_text.strip():
        raise HTTPException(status_code=422, detail="OCR produced no text.")

    page_count = max(1, ocr_text.count("--- Page "))
    parsed = call_llm(ocr_text)

    def s(key, n=300): return str(parsed.get(key, "")).strip()[:n]

    return InvoiceData(
        invoice_number=s("invoice_number", 60),
        eway_bill_number=s("eway_bill_number", 30),
        invoice_date=s("invoice_date", 30),
        mode_of_payment=s("mode_of_payment", 60),
        supplier_ref=s("supplier_ref", 60),
        buyer_order_number=s("buyer_order_number", 60),
        dispatch_document_number=s("dispatch_document_number", 60),
        dispatched_through=s("dispatched_through", 100),
        destination=s("destination", 100),
        delivery_note=s("delivery_note", 60),

        vendor_name=s("vendor_name", 150),
        vendor_address=s("vendor_address", 300),
        vendor_gstin=s("vendor_gstin", 20),
        vendor_state=s("vendor_state", 100),
        vendor_email=s("vendor_email", 100),

        buyer_name=s("buyer_name", 150),
        buyer_address=s("buyer_address", 300),
        buyer_gstin=s("buyer_gstin", 20),
        buyer_state=s("buyer_state", 100),

        items=[
            LineItem(
                sl_no=str(i.get("sl_no", ""))[:10],
                description=str(i.get("description", ""))[:200],
                batch=str(i.get("batch", ""))[:50],
                hsn_sac=str(i.get("hsn_sac", ""))[:20],
                quantity=str(i.get("quantity", ""))[:30],
                rate=str(i.get("rate", ""))[:30],
                per=str(i.get("per", ""))[:20],
                amount=str(i.get("amount", ""))[:30],
            )
            for i in parsed.get("items", []) if isinstance(i, dict)
        ],

        taxable_value=s("taxable_value", 30),
        cgst_rate=s("cgst_rate", 10),
        cgst_amount=s("cgst_amount", 30),
        sgst_rate=s("sgst_rate", 10),
        sgst_amount=s("sgst_amount", 30),
        igst_rate=s("igst_rate", 10),
        igst_amount=s("igst_amount", 30),
        output_cgst=s("output_cgst", 30),
        output_sgst=s("output_sgst", 30),
        total_tax_amount=s("total_tax_amount", 30),
        grand_total=s("grand_total", 30),
        amount_in_words=s("amount_in_words", 300),
        tax_amount_in_words=s("tax_amount_in_words", 300),

        hsn_summary=[
            HSNSummary(
                hsn_sac=str(h.get("hsn_sac", ""))[:20],
                taxable_value=str(h.get("taxable_value", ""))[:30],
                cgst_rate=str(h.get("cgst_rate", ""))[:10],
                cgst_amount=str(h.get("cgst_amount", ""))[:30],
                sgst_rate=str(h.get("sgst_rate", ""))[:10],
                sgst_amount=str(h.get("sgst_amount", ""))[:30],
                total_tax=str(h.get("total_tax", ""))[:30],
            )
            for h in parsed.get("hsn_summary", []) if isinstance(h, dict)
        ],

        declaration=s("declaration", 500),
        authorised_signatory=s("authorised_signatory", 100),
        is_computer_generated=bool(parsed.get("is_computer_generated", False)),
        source_pages=page_count,
    )


@app.get("/health")
async def health():
    return {"status": "healthy", "model": LLM_MODEL}


HTML_PATH = Path(__file__).parent / "index.html"

@app.get("/", response_class=HTMLResponse)
async def serve_ui():
    if not HTML_PATH.exists():
        return HTMLResponse("<h2>index.html not found</h2>", 500)
    return HTMLResponse(HTML_PATH.read_text(encoding="utf-8"))


if __name__ == "__main__":
    import uvicorn
    port = int(__import__("os").environ.get("HF_PORT", 7860))
    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)