core/extract.py

"""Receipt + payment extraction (MiniCPM-V-4.6, via core.inference).

v2 data model — a transaction is flexible, so real bills fit without hardcoding:
{
  "vendor": str, "date": "YYYY-MM-DD"|"", "currency": str,
  "line_items": [ {"name": str, "qty": number, "amount": number, "category": str} ],
  "charges":    [ {"label": str, "amount": number} ],   # ANY taxes/fees/discount/round-off
  "total": number, "category": str, "note": str, "source": str
}
Reconciliation: sum(line_items.amount) + sum(charges.amount) ≈ total.
(Discounts and round-downs are negative charge amounts.)
"""

from __future__ import annotations

import json
import re
from typing import Any

from core.analytics import parse_date as _parse_date
from core import inference

gpu_decorator = inference.gpu_decorator
MAX_NEW_TOKENS = 1024

# --------------------------------------------------------------------------- #
# Prompts
# --------------------------------------------------------------------------- #
SYSTEM_PROMPT = (
    "You are an expert receipt and bill reader for a budgeting app. Real bills "
    "are messy and varied: missing totals, several tax lines (e.g. SGST and CGST "
    "separately), service charges, tips, discounts, manual round-off, and mixed "
    "items. Read EVERYTHING carefully and miss nothing.\n"
    "Return ONLY one valid JSON object — no prose, no markdown — EXACTLY:\n"
    "{\n"
    '  "vendor": string,            // merchant/store name, "" if unknown\n'
    '  "date": string,              // bill date as YYYY-MM-DD (convert any format), "" if none\n'
    '  "currency": string,          // code or symbol seen, e.g. "INR","₹","USD","$"\n'
    '  "line_items": [              // EVERY purchased item row, in order\n'
    '    { "name": string, "qty": number, "amount": number }   // amount = line total\n'
    "  ],\n"
    '  "charges": [                 // EVERY non-item money line, each separately\n'
    '    { "label": string, "amount": number }   // e.g. {"label":"SGST 9%","amount":31.95}\n'
    "  ],\n"
    '  "total": number              // grand total actually payable\n'
    "}\n"
    "Rules:\n"
    "- line_items are the REAL things bought. Do NOT put subtotal, total, grand "
    "total, net amount, amount payable, balance, change/tendered, or 'amount in "
    "words' in line_items — those summarise the bill; the payable one is `total`.\n"
    "- Never repeat the same row twice. Each item appears once.\n"
    "- Capture EACH tax line as its OWN charges entry (do NOT merge SGST + CGST).\n"
    "- Service charge, tip → positive charges. Discount, round-down → NEGATIVE amount.\n"
    "- A round-off that lowers the total is a charge like {\"label\":\"Round off\",\"amount\":-0.40}.\n"
    "- Utility/telecom bills: the CONSUMPTION lines (energy charges, fixed/demand "
    "charges, fuel surcharge, data/talktime) ARE line_items. Only government taxes/"
    "duties (electricity duty, GST) go in charges.\n"
    "- Numbers are plain (no symbols/commas), dot decimal. qty defaults to 1.\n"
    "- If no total is printed, COMPUTE it = sum(line_items) + sum(charges).\n"
    "- Check: sum(line_items) + sum(charges) must equal total. If not, re-examine "
    "ONCE for missed items, missed tax lines, or misplaced decimals, then fix.\n"
    "Output the JSON object and nothing else."
)

USER_PROMPT = "Extract this receipt/bill into the required JSON. Return ONLY the JSON object."

RETRY_PROMPT = (
    "Your previous answer was not valid JSON in the required schema. Return ONLY a "
    "single JSON object with keys: vendor (string), date (string), currency (string), "
    "line_items (array of {name, qty, amount}), charges (array of {label, amount}), "
    "total (number). Capture every item and every tax line separately. No markdown."
)

PAYMENT_SYSTEM_PROMPT = (
    "You read a screenshot of one digital payment (UPI, GPay, PhonePe, card, or a "
    "bank app). Extract that single transaction. Return ONLY one JSON object — no "
    "prose, no markdown — EXACTLY:\n"
    "{\n"
    '  "vendor": string,   // payee / merchant / person paid, "" if unknown\n'
    '  "date": string,     // payment date as YYYY-MM-DD, "" if none\n'
    '  "currency": string, // code or symbol, e.g. "INR","₹","$"\n'
    '  "amount": number,   // amount paid (positive plain number)\n'
    '  "note": string      // any reference/description shown, "" if none\n'
    "}\n"
    "Numbers are plain (no symbols/commas), dot decimal. Output only the JSON."
)
PAYMENT_USER_PROMPT = "Extract this payment screenshot into the required JSON. Return ONLY the JSON object."

EMPTY_RESULT: dict[str, Any] = {
    "vendor": "", "date": "", "currency": "",
    "line_items": [], "charges": [], "total": 0,
}

RECON_ABS_TOL = 0.5
RECON_PCT_TOL = 0.01


# --------------------------------------------------------------------------- #
# Parsing helpers
# --------------------------------------------------------------------------- #
_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.DOTALL | re.IGNORECASE)


def _strip_fences(text: str) -> str:
    m = _FENCE_RE.search(text)
    return m.group(1).strip() if m else text.strip()


def _extract_json_object(text: str) -> str:
    start = text.find("{")
    if start == -1:
        return text
    depth = 0
    for i in range(start, len(text)):
        if text[i] == "{":
            depth += 1
        elif text[i] == "}":
            depth -= 1
            if depth == 0:
                return text[start : i + 1]
    return text[start:]


def _coerce_number(value: Any) -> float:
    if isinstance(value, (int, float)):
        return float(value)
    if isinstance(value, str):
        cleaned = re.sub(r"[^0-9.\-]", "", value.replace(",", ""))
        try:
            return float(cleaned) if cleaned not in ("", "-", ".") else 0.0
        except ValueError:
            return 0.0
    return 0.0


def _normalize_date(value: Any) -> str:
    d = _parse_date(value)
    return d.isoformat() if d else ""


def _coerce_charge(ch: Any) -> dict[str, Any] | None:
    """Accept {'label','amount'} or {'<label>': <amount>}; return {label, amount}."""
    if not isinstance(ch, dict):
        return None
    if "amount" in ch or "label" in ch:
        return {"label": str(ch.get("label", "") or "Charge"),
                "amount": _coerce_number(ch.get("amount", 0))}
    # single-pair form
    for k, v in ch.items():
        return {"label": str(k), "amount": _coerce_number(v)}
    return None


def _validate(data: Any) -> dict[str, Any]:
    if not isinstance(data, dict):
        raise ValueError("Top-level JSON is not an object")
    result: dict[str, Any] = {
        "vendor": str(data.get("vendor", "") or ""),
        "date": _normalize_date(data.get("date", "")),
        "currency": str(data.get("currency", "") or ""),
        "line_items": [],
        "charges": [],
        "total": _coerce_number(data.get("total", 0)),
    }
    for it in data.get("line_items", []) or []:
        if isinstance(it, dict):
            result["line_items"].append({
                "name": str(it.get("name", "") or ""),
                "qty": _coerce_number(it.get("qty", 1)) or 1,
                "amount": _coerce_number(it.get("amount", 0)),
            })
    for ch in data.get("charges", []) or []:
        c = _coerce_charge(ch)
        if c and (c["label"] or c["amount"]):
            result["charges"].append(c)
    return result


# --------------------------------------------------------------------------- #
# Deterministic cleanup — understand a literal extraction (no model, instant)
# --------------------------------------------------------------------------- #
# A row whose NAME is a bill summary, not a purchased item (drop from line_items;
# its value is a candidate grand total). Matches names made up ENTIRELY of summary
# words, so multi-word labels ("Total Amount Payable") are caught while brand/item
# names with other words ("Total Wireless Recharge") are NOT.
_SUMMARY_RE = re.compile(
    r"^\s*(?:(?:sub|grand|final|net|nett|gross|total|amount|payable|balance|due|"
    r"value|bill|invoice|qty|items?|to|pay|in|words|tendered|change|cash|paid)"
    r"\b[\s:./\-]*)+$", re.I)

# A 'charge'-named line that is actually a real consumption ITEM (utility/telecom).
# Checked BEFORE the tax/fee test so "energy charges" stays an item.
_ITEM_KEEP_RE = re.compile(
    r"\b(energy|fixed|demand|wheeling|meter|consumption|sanction|connection|"
    r"fuel\s*(?:adj\w*|surcharge|charge)|fppca|fpppa|\bfac\b|"
    r"data|talktime|talk\s*time|usage|rental|plan|pack|recharge)\b", re.I)

# A line that belongs in `charges` (tax / statutory / service fee), never an item.
_TAXFEE_RE = re.compile(
    r"\b(c?gst|sgst|igst|utgst|\bvat\b|sales\s*tax|service\s*tax|service\s*charge|"
    r"svc\s*charge|\bcess\b|\btcs\b|\btds\b|octroi|electricity\s*duty|\bduty\b|"
    r"delivery|packing|packag\w*|conveni\w*|handling|platform\s*fee|"
    r"gratuity|\btip\b|round[\s-]*off|rounding|\bdiscount\b|cashback|coupon|"
    r"\bsavings?\b|\bsurcharge\b|\bfee\b|\btax\b)\b", re.I)


def _row_kind(name: str) -> str:
    n = name or ""
    if _SUMMARY_RE.match(n):
        return "summary"
    if _ITEM_KEEP_RE.search(n):
        return "item"
    if _TAXFEE_RE.search(n):
        return "charge"
    return "item"


def clean_extraction(record: dict[str, Any]) -> dict[str, Any]:
    """Turn a literal OCR extraction into a coherent transaction — instantly, with
    no model call. Drops summary/total rows mis-read as items, moves tax/fee rows
    out of line_items into charges (while keeping utility consumption rows as
    items), de-duplicates, and recovers a sensible total. Idempotent."""
    rec = dict(record)
    items_in = rec.get("line_items") or []
    charges_in = rec.get("charges") or []

    kept_items: list[dict[str, Any]] = []
    moved_charges: list[dict[str, Any]] = []
    clean_charges: list[dict[str, Any]] = []
    total_candidates: list[float] = []
    seen: set[tuple[str, float]] = set()

    for it in items_in:
        name = str(it.get("name", "") or "").strip()
        amount = _coerce_number(it.get("amount", 0))
        kind = _row_kind(name)
        if kind == "summary":
            if amount:
                total_candidates.append(amount)
            continue
        if kind == "charge":
            moved_charges.append({"label": name or "Charge", "amount": amount})
            continue
        key = (name.lower(), round(amount, 2))
        if name and key in seen:   # drop an exact duplicate item row
            continue
        seen.add(key)
        kept_items.append(it)

    # Remove summary rows that landed in charges; keep genuine charges.
    for c in charges_in:
        label = str(c.get("label", "") or "").strip()
        amount = _coerce_number(c.get("amount", 0))
        if _SUMMARY_RE.match(label):
            if amount:
                total_candidates.append(amount)
            continue
        clean_charges.append({"label": label or "Charge", "amount": amount})
    clean_charges.extend(moved_charges)

    rec["line_items"] = kept_items
    rec["charges"] = clean_charges

    printed = _coerce_number(rec.get("total", 0))
    computed = round(sum(_coerce_number(i.get("amount", 0)) for i in kept_items)
                     + sum(_coerce_number(c.get("amount", 0)) for c in clean_charges), 2)
    if printed <= 0:
        # Prefer a printed grand-total we pulled off a summary row; else compute.
        rec["total"] = max(total_candidates) if total_candidates else computed
    return rec


# --------------------------------------------------------------------------- #
# Reconciliation
# --------------------------------------------------------------------------- #
def reconcile(record: dict[str, Any]) -> dict[str, Any]:
    items = record.get("line_items") or []
    total = _coerce_number(record.get("total", 0))
    if not items:
        return {"items_sum": 0.0, "charges_sum": 0.0, "expected_total": round(total, 2),
                "total": round(total, 2), "gap": 0.0, "tolerance": 0.0, "ok": True,
                "message": "No line items to reconcile."}
    items_sum = sum(_coerce_number(it.get("amount", 0)) for it in items)
    charges_sum = sum(_coerce_number(c.get("amount", 0)) for c in (record.get("charges") or []))
    expected = round(items_sum + charges_sum, 2)
    gap = round(total - expected, 2)
    tolerance = max(RECON_ABS_TOL, RECON_PCT_TOL * abs(total))
    ok = abs(gap) <= tolerance
    message = (f"Items + charges = {expected:.2f} ≈ total {total:.2f}." if ok
               else f"Items + charges = {expected:.2f}, but total reads {total:.2f} "
                    f"(off by {gap:+.2f}).")
    return {"items_sum": round(items_sum, 2), "charges_sum": round(charges_sum, 2),
            "expected_total": expected, "total": total, "gap": gap,
            "tolerance": round(tolerance, 2), "ok": ok, "message": message}


def parse_receipt_json(text: str) -> dict[str, Any]:
    return _validate(json.loads(_extract_json_object(_strip_fences(text))))


def parse_payment_json(text: str) -> dict[str, Any]:
    data = json.loads(_extract_json_object(_strip_fences(text)))
    if not isinstance(data, dict):
        raise ValueError("Top-level JSON is not an object")
    rec = dict(EMPTY_RESULT)
    rec["line_items"], rec["charges"] = [], []
    rec["vendor"] = str(data.get("vendor", "") or "")
    rec["date"] = _normalize_date(data.get("date", ""))
    rec["currency"] = str(data.get("currency", "") or "")
    rec["total"] = _coerce_number(data.get("amount", 0))
    rec["note"] = str(data.get("note", "") or "")
    rec["source"] = "payment"
    return rec


# --------------------------------------------------------------------------- #
# Inference (vision, via core.inference)
# --------------------------------------------------------------------------- #
def _run_model(image, prompt: str, system: str = SYSTEM_PROMPT) -> str:
    return inference.vision_generate(image, system, prompt, MAX_NEW_TOKENS)


def extract_receipt(image) -> dict[str, Any]:
    if image is None:
        return {**EMPTY_RESULT, "line_items": [], "charges": [], "_error": "No image provided."}
    for prompt in (USER_PROMPT, RETRY_PROMPT):
        try:
            return clean_extraction(parse_receipt_json(_run_model(image, prompt)))
        except Exception as e:
            print(f"[extract] receipt parse failed: {e}")
    return {**EMPTY_RESULT, "line_items": [], "charges": [],
            "_error": "Could not read the receipt clearly — please edit below."}


def extract_payment(image) -> dict[str, Any]:
    if image is None:
        return {**EMPTY_RESULT, "line_items": [], "charges": [], "source": "payment",
                "note": "", "_error": "No image provided."}
    for _ in range(2):
        try:
            return parse_payment_json(_run_model(image, PAYMENT_USER_PROMPT, PAYMENT_SYSTEM_PROMPT))
        except Exception as e:
            print(f"[extract] payment parse failed: {e}")
    return {**EMPTY_RESULT, "line_items": [], "charges": [], "source": "payment",
            "note": "", "_error": "Could not read the payment screenshot."}