anchors.py

"""Deterministic regex anchors + merge/validation for the known invoice template.

No heavy imports here so it can be unit-tested locally without torch/gradio.
"""
import json
import os
import re

try:
    VENDORS = json.loads(os.environ.get("VENDOR_RECORDS", "[]"))
except json.JSONDecodeError:
    VENDORS = []

GSTIN_RE = re.compile(r"\b\d{2}[A-Z]{5}\d{4}[A-Z][0-9A-Z]Z[0-9A-Z]\b")
PAN_RE = re.compile(r"\b[A-Z]{5}\d{4}[A-Z]\b")
_AMT = r"([\d,]+(?:\.\d+)?)"


def _num(s):
    try:
        n = float(s.replace(",", ""))
        return int(n) if n.is_integer() else n
    except (ValueError, AttributeError):
        return None


def _search(pattern, text, flags=0):
    m = re.search(pattern, text, flags)
    return m.group(1).strip() if m else None


_NUM_TOKEN = re.compile(r"^\d[\d,]*(?:\.\d+)?$")


def parse_annexure(text: str, grand_total) -> list | None:
    """Parse the employee annexure table by its numeric column structure.

    Rows look like: <manager> <employee...> <13-15 numbers> [remark]
    Common columns: idx0 = monthly billing, idx9 = total payable days,
    idx10 = payable billing; the tail differs per template:
      13 numbers -> billing, ..., payable, charges, total
      14 numbers -> billing, ..., payable, cgst, sgst, total
      15 numbers -> billing, ..., payable, charges, subtotal, gst, total
    Returns rows only when their totals reconcile with the invoice grand total.
    """
    m = re.search(r"Kind Attention Person.*?\n", text)
    if not m:
        return None
    rows = []
    for line in text[m.end():].splitlines():
        tokens = line.split()
        if not tokens or tokens[0].upper() == "TOTAL":
            continue
        nums, name_tokens, remark_tokens = [], [], []
        for t in tokens:
            cleaned = t.strip()
            if _NUM_TOKEN.match(cleaned):
                nums.append(_num(cleaned))
            elif not nums:
                name_tokens.append(t)
            else:
                remark_tokens.append(t)
        if len(nums) < 13 or len(nums) > 15:
            continue
        row = {
            "name": " ".join(name_tokens),
            "monthly_billing": nums[0],
            "payable_days": nums[9],
            "amount": nums[10],
            "gst_amount": None,
            "total": nums[-1],
        }
        if len(nums) == 13:
            row["charges"] = nums[11]
        elif len(nums) == 14:
            row["gst_amount"] = (nums[11] or 0) + (nums[12] or 0)
        elif len(nums) == 15:
            row["charges"] = nums[11]
            row["gst_amount"] = nums[13]
        if remark_tokens:
            row["remark"] = " ".join(remark_tokens)
        rows.append(row)
    if not rows:
        return None
    if isinstance(grand_total, (int, float)):
        if abs(sum(r["total"] or 0 for r in rows) - grand_total) > 1:
            return None  # structure didn't match; let the model's rows stand
    return rows


def deterministic_fields(text: str) -> dict:
    """Extract label-anchored fields with regex; reliable on the known template."""
    f = {}
    if v := _search(r"Bill No\.?\s*:?\s*([A-Z]{2,6}/\d+/\d{2}-\d{2})", text):
        f["invoice_number"] = v
    if v := _search(r"Bill Date\s*:?\s*(\d{2}-\d{2}-\d{2,4})", text):
        f["invoice_date"] = v
    if v := _search(r"Due Date\s*:?\s*(\d{2}-\d{2}-\d{2,4})", text):
        f["due_date"] = v
    if v := _search(r"SAC Code\s*:?\s*(\d{4,6})", text):
        f["sac_code"] = v
    if v := _search(r"(?:WO No\.?\s*:?-?\s*|Work Order\s*:?\s*)([A-Z0-9/\-]{4,})", text):
        f["work_order"] = v
    if v := _search(r"Rupees\s*:\s*(.+?Only)", text, re.S):
        f["amount_in_words"] = re.sub(r"\s+", " ", v)

    grand = _num(_search(r"Rupees\s*:.+?Only\s+" + _AMT, text, re.S))
    if grand is not None:
        # When the total is anchored, the whole amounts object is authoritative:
        # a tax type with no "Add : XGST" line was not charged (null), so model
        # hallucinations can never leak through the merge.
        amounts = {
            "taxable_value": None,
            "cgst": _num(_search(r"CGST\s*@\s*\d+%\s+" + _AMT, text)),
            "sgst": _num(_search(r"SGST\s*@\s*\d+%\s+" + _AMT, text)),
            "igst": _num(_search(r"IGST\s*@\s*\d+%\s+" + _AMT, text)),
            "grand_total": grand,
        }
        taxes = sum(amounts[k] or 0 for k in ("cgst", "sgst", "igst"))
        tv = round(grand - taxes, 2)
        amounts["taxable_value"] = int(tv) if tv == int(tv) else tv
        f["amounts"] = amounts

    if (rows := parse_annexure(text, grand)) is not None:
        f["employees"] = rows

    gstins = list(dict.fromkeys(GSTIN_RE.findall(text)))
    pans = list(dict.fromkeys(PAN_RE.findall(text)))
    vendor = next(
        (
            v for v in VENDORS
            if (v.get("gstin") in gstins)
            or (v.get("prefix") and f.get("invoice_number", "").startswith(v["prefix"] + "/"))
        ),
        None,
    )
    if vendor:
        f["vendor"] = {
            "name": vendor.get("name"),
            "address": vendor.get("address"),
            "gstin": vendor.get("gstin"),
            "pan": vendor.get("pan"),
            "email": vendor.get("email"),
            "phone": vendor.get("phone"),
            "pf_no": vendor.get("pf_no"),
            "esic_no": vendor.get("esic_no"),
        }
        if vendor.get("bank"):
            f["bank_details"] = vendor["bank"]
        # On the known template, a missing WO/Work Order regex hit means there is
        # none - don't let a model guess leak through.
        f.setdefault("work_order", None)
        buyer = {}
        # Buyer name = the addressee line right after "Original For Recepient",
        # with the right-hand column ("Bill No. : ...") stripped off.
        if m := re.search(r"Original For Recepient\s*\n(.+)", text):
            name = re.split(r"\s{2,}|Bill No", m.group(1), maxsplit=1)[0].strip()
            if len(name) > 3:
                buyer["name"] = name
        if g := next((g for g in gstins if g != vendor.get("gstin")), None):
            buyer["gstin"] = g
        vendor_pan = vendor.get("pan")
        buyer_gstin_pan = buyer.get("gstin", "")[2:12] or None
        if p := next((p for p in pans if p != vendor_pan), None):
            # prefer the PAN embedded in the buyer's GSTIN when available
            buyer["pan"] = buyer_gstin_pan if buyer_gstin_pan in pans else p
        elif buyer_gstin_pan:
            buyer["pan"] = buyer_gstin_pan
        if buyer:
            f["buyer"] = buyer
    return f


# Anchored versions of these keys are complete and authoritative - replace the
# model's value entirely (a merged-in None means "verified absent", not unknown).
_REPLACE_KEYS = {"amounts", "vendor", "employees", "bank_details"}


def merge_result(llm_result: dict, anchors: dict) -> dict:
    """Anchored fields override the model output; nested dicts merge per-key."""
    out = dict(llm_result) if isinstance(llm_result, dict) else {}
    for k, v in anchors.items():
        if isinstance(v, dict) and k not in _REPLACE_KEYS:
            base = out.get(k) if isinstance(out.get(k), dict) else {}
            out[k] = {**base, **{kk: vv for kk, vv in v.items() if vv is not None}}
        else:
            out[k] = v

    warnings = []
    a = out.get("amounts") or {}
    gt, tv = a.get("grand_total"), a.get("taxable_value")
    taxes = sum(a.get(k) or 0 for k in ("cgst", "sgst", "igst"))
    if isinstance(gt, (int, float)) and isinstance(tv, (int, float)):
        if abs((tv + taxes) - gt) > 1:
            warnings.append("amounts do not reconcile: taxable + taxes != grand_total")
    else:
        warnings.append("could not verify amounts against grand total")
    emp_totals = [
        e.get("total")
        for e in (out.get("employees") or [])
        if isinstance(e, dict) and isinstance(e.get("total"), (int, float))
    ]
    if emp_totals and isinstance(gt, (int, float)) and abs(sum(emp_totals) - gt) > 1:
        warnings.append("employee annexure rows may be unreliable: row totals != grand_total")
    if warnings:
        out["warnings"] = warnings
    return out