extraction.py

"""
DocuMint AI - Core Invoice/Document Extraction Engine

Supports multiple extraction backends:
1. Hugging Face Inference API (Qwen2.5-VL) - Primary, highest accuracy
2. Tesseract OCR + regex patterns - Fallback, CPU-only, zero cost
"""

import os
import re
import json
import base64
import logging
from io import BytesIO
from typing import Optional
from dataclasses import dataclass, field, asdict
from datetime import datetime

import requests
from PIL import Image
from huggingface_hub import InferenceClient

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Data Models
# ---------------------------------------------------------------------------

@dataclass
class LineItem:
    description: str = ""
    quantity: Optional[float] = None
    unit_price: Optional[float] = None
    amount: Optional[float] = None

@dataclass
class BoundingBox:
    field: str = ""
    x1: int = 0
    y1: int = 0
    x2: int = 0
    y2: int = 0
    page: int = 0

@dataclass
class ExtractionResult:
    # Legacy named fields (populated for invoice use case, empty for others)
    vendor_name: str = ""
    vendor_address: str = ""
    invoice_number: str = ""
    invoice_date: str = ""
    due_date: str = ""
    currency: str = ""
    subtotal: Optional[float] = None
    tax_amount: Optional[float] = None
    tax_rate: Optional[str] = None
    total_amount: Optional[float] = None
    line_items: list = field(default_factory=list)
    payment_terms: str = ""
    customer_name: str = ""
    customer_address: str = ""
    raw_text: str = ""
    extraction_method: str = ""
    confidence: float = 0.0
    processing_time_ms: int = 0
    bounding_boxes: list = field(default_factory=list)
    page_image: Optional[str] = None  # data-URI of rendered page (for PDFs)
    # Dynamic fields — primary data carrier for all use cases
    document_type: str = ""
    extracted_fields: dict = field(default_factory=dict)

    def to_dict(self):
        d = asdict(self)
        d["line_items"] = [asdict(li) if isinstance(li, LineItem) else li for li in self.line_items]
        d["bounding_boxes"] = [asdict(bb) if isinstance(bb, BoundingBox) else bb for bb in self.bounding_boxes]
        return d


# ---------------------------------------------------------------------------
# Bounding Box Generation - matches extracted values to OCR word positions
# ---------------------------------------------------------------------------

def _get_ocr_word_boxes(image: Image.Image) -> list[dict]:
    """Run pytesseract image_to_data to get word-level bounding boxes."""
    try:
        import pytesseract
        data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT, config="--psm 3")
        words = []
        for i in range(len(data["text"])):
            text = data["text"][i].strip()
            conf = int(data["conf"][i])
            if text and conf > 10:
                words.append({
                    "text": text,
                    "x": data["left"][i],
                    "y": data["top"][i],
                    "w": data["width"][i],
                    "h": data["height"][i],
                })
        return words
    except Exception as e:
        logger.warning(f"OCR word boxes failed: {e}")
        return []


def _date_search_variants(date_str: str) -> list[str]:
    """Generate alternative date representations for OCR matching.

    VLM returns YYYY-MM-DD but documents may show DD-Mon-YYYY, DD/MM/YYYY, etc.
    """
    variants = [date_str]
    # Try parsing YYYY-MM-DD
    m = re.match(r"(\d{4})-(\d{2})-(\d{2})", date_str)
    if m:
        y, mo, d = m.group(1), m.group(2), m.group(3)
        d_int = str(int(d))  # strip leading zero
        months = ["Jan","Feb","Mar","Apr","May","Jun","Jul","Aug","Sep","Oct","Nov","Dec"]
        month_name = months[int(mo) - 1] if 1 <= int(mo) <= 12 else mo
        month_full = ["January","February","March","April","May","June",
                      "July","August","September","October","November","December"]
        full_name = month_full[int(mo) - 1] if 1 <= int(mo) <= 12 else mo
        variants += [
            f"{d}-{month_name}-{y}",       # 15-Aug-2026
            f"{d_int}-{month_name}-{y}",    # 5-Aug-2026
            f"{d}/{mo}/{y}",                # 15/08/2026
            f"{d_int}/{mo}/{y}",            # 5/08/2026
            f"{mo}/{d}/{y}",                # 08/15/2026
            f"{d}.{mo}.{y}",                # 15.08.2026
            f"{d}-{mo}-{y}",                # 15-08-2026
            f"{d_int} {month_name} {y}",    # 5 Aug 2026
            f"{d} {month_name} {y}",        # 15 Aug 2026
            f"{month_name} {d_int}, {y}",   # Aug 5, 2026
            f"{d_int} {full_name} {y}",     # 5 August 2026
            f"{d} {full_name} {y}",         # 15 August 2026
            f"{full_name} {d_int}, {y}",    # August 5, 2026
        ]
    return list(dict.fromkeys(variants))


def _currency_search_variants(currency: str) -> list[str]:
    """Generate alternative currency representations for OCR matching."""
    symbol_map = {
        "INR": ["₹", "Rs", "Rs.", "INR"],
        "USD": ["$", "USD", "US$"],
        "EUR": ["€", "EUR"],
        "GBP": ["£", "GBP"],
        "TRY": ["₺", "TL", "TRY"],
        "JPY": ["¥", "JPY"],
    }
    return symbol_map.get(currency.upper(), [currency])


def _word_matches_token(word_text: str, token: str) -> bool:
    """Check if an OCR word matches a search token.

    Uses full-word matching for alphabetic tokens (avoiding "tax" matching
    "taxation") and substring matching for numeric tokens (since OCR often
    merges numbers with punctuation like "$1,234.56").
    """
    wt = word_text.lower()
    # Numeric tokens: allow substring (OCR may merge currency symbols, commas)
    # Also strip common OCR artifacts before checking
    if any(c.isdigit() for c in token):
        # Direct substring
        if token in wt:
            return True
        # Strip non-digit non-dot non-comma chars and retry
        wt_digits = re.sub(r"[^0-9.,]", "", wt)
        tk_digits = re.sub(r"[^0-9.,]", "", token)
        return tk_digits != "" and tk_digits in wt_digits
    # Short tokens (<=2 chars): exact match only to avoid noise
    if len(token) <= 2:
        return wt == token
    # Alphabetic tokens: check word boundaries — strip common punctuation
    # from OCR word and compare
    wt_clean = re.sub(r"[^a-z0-9]", "", wt)
    tk_clean = re.sub(r"[^a-z0-9]", "", token)
    if not tk_clean:
        return False
    return wt_clean == tk_clean or wt_clean.startswith(tk_clean) or wt_clean.endswith(tk_clean)


def _find_all_text_bboxes(words: list[dict], search_text: str) -> list[BoundingBox]:
    """Find ALL bounding boxes for a text value by matching against OCR words.

    Returns every match location so callers can disambiguate using context
    (e.g. proximity to a field label).
    """
    if not search_text or not words:
        return []

    search_lower = search_text.lower().strip()
    search_tokens = search_lower.split()

    if not search_tokens:
        return []

    results = []

    if len(search_tokens) == 1:
        token = search_tokens[0]
        for w in words:
            if _word_matches_token(w["text"], token):
                results.append(BoundingBox(
                    x1=w["x"], y1=w["y"],
                    x2=w["x"] + w["w"], y2=w["y"] + w["h"],
                ))
        # Fallback: try joining the search text without spaces (OCR may merge)
        if not results:
            merged = search_lower.replace(" ", "")
            if len(merged) >= 3:
                for w in words:
                    wt_clean = re.sub(r"[^a-z0-9]", "", w["text"].lower())
                    if wt_clean == merged:
                        results.append(BoundingBox(
                            x1=w["x"], y1=w["y"],
                            x2=w["x"] + w["w"], y2=w["y"] + w["h"],
                        ))
    else:
        # Multi-word: try anchoring on EACH token position, not just the first.
        # OCR may miss or garble the first word but read later words fine.
        for anchor_idx in range(len(search_tokens)):
            anchor_token = search_tokens[anchor_idx]
            for i, w in enumerate(words):
                if not _word_matches_token(w["text"], anchor_token):
                    continue
                # Expand backward to find earlier tokens
                x1, y1 = w["x"], w["y"]
                x2, y2 = w["x"] + w["w"], w["y"] + w["h"]
                matched = 1

                # Look backward for tokens before the anchor
                if anchor_idx > 0:
                    back_j = anchor_idx - 1
                    for bk in range(1, min(anchor_idx * 2 + 1, i + 1)):
                        if back_j < 0:
                            break
                        prev_w = words[i - bk]
                        if abs(prev_w["y"] - w["y"]) > 80:
                            break
                        if _word_matches_token(prev_w["text"], search_tokens[back_j]):
                            x1 = min(x1, prev_w["x"])
                            y1 = min(y1, prev_w["y"])
                            x2 = max(x2, prev_w["x"] + prev_w["w"])
                            y2 = max(y2, prev_w["y"] + prev_w["h"])
                            matched += 1
                            back_j -= 1

                # Look forward for tokens after the anchor
                search_j = anchor_idx + 1
                scan_limit = min((len(search_tokens) - anchor_idx) * 2, len(words) - i)
                for k in range(1, scan_limit):
                    if search_j >= len(search_tokens):
                        break
                    next_w = words[i + k]
                    if abs(next_w["y"] - w["y"]) > 80:
                        break
                    if _word_matches_token(next_w["text"], search_tokens[search_j]):
                        x2 = max(x2, next_w["x"] + next_w["w"])
                        y2 = max(y2, next_w["y"] + next_w["h"])
                        x1 = min(x1, next_w["x"])
                        y1 = min(y1, next_w["y"])
                        matched += 1
                        search_j += 1

                # Require at least half the tokens to match
                threshold = max(1, (len(search_tokens) + 1) // 2)
                if matched >= threshold:
                    results.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2))

            # If we found matches anchoring on this token, don't try other anchors
            # (avoids duplicate/overlapping results)
            if results:
                break

    return results


def _find_text_bbox(words: list[dict], search_text: str) -> Optional[BoundingBox]:
    """Find bounding box for a text value (first match). Legacy convenience wrapper."""
    matches = _find_all_text_bboxes(words, search_text)
    return matches[0] if matches else None


def _bbox_center(bb: BoundingBox) -> tuple[float, float]:
    """Return (cx, cy) center of a bounding box."""
    return ((bb.x1 + bb.x2) / 2, (bb.y1 + bb.y2) / 2)


def _bbox_distance(a: BoundingBox, b: BoundingBox) -> float:
    """Euclidean distance between bounding box centers."""
    ax, ay = _bbox_center(a)
    bx, by = _bbox_center(b)
    return ((ax - bx) ** 2 + (ay - by) ** 2) ** 0.5


def _bbox_overlaps(a: BoundingBox, b: BoundingBox, margin: int = 5) -> bool:
    """Check if two bounding boxes overlap (with optional margin)."""
    return not (
        a.x2 + margin < b.x1 or b.x2 + margin < a.x1 or
        a.y2 + margin < b.y1 or b.y2 + margin < a.y1
    )


def _field_name_to_label_tokens(field_name: str) -> list[str]:
    """Derive likely label keywords from a snake_case field name.

    e.g. 'invoice_date' → ['invoice', 'date']
         'total_amount' → ['total', 'amount']
         'pan'          → ['pan']
    """
    return [t.lower() for t in field_name.split("_") if len(t) > 1]


def _find_label_bboxes(words: list[dict], field_name: str) -> list[BoundingBox]:
    """Find ALL positions where a field's label appears in the document OCR.

    Returns all matches so the caller can pick the one closest to the value.
    Searches for the label keywords derived from the field name. Looks for
    multi-word labels first (e.g. "Invoice Date"), then single keywords.
    """
    tokens = _field_name_to_label_tokens(field_name)
    if not tokens:
        return []

    # Try the full label as a phrase first
    full_label = " ".join(tokens)
    matches = _find_all_text_bboxes(words, full_label)
    if matches:
        return matches

    # Fall back to the most distinctive single token (longest word)
    tokens_by_len = sorted(tokens, key=len, reverse=True)
    for token in tokens_by_len:
        # Skip very common short words that appear everywhere
        if token in ("of", "no", "in", "to", "at", "by", "id", "or"):
            continue
        matches = _find_all_text_bboxes(words, token)
        if matches:
            return matches

    return []


def _find_label_bbox(words: list[dict], field_name: str) -> Optional[BoundingBox]:
    """Find where a field's label appears — convenience wrapper returning first match."""
    matches = _find_label_bboxes(words, field_name)
    return matches[0] if matches else None


def _pick_best_match(
    candidates: list[BoundingBox],
    label_bbox: Optional[BoundingBox],
    claimed: list[BoundingBox],
    label_bboxes: Optional[list[BoundingBox]] = None,
) -> Optional[BoundingBox]:
    """Pick the best bbox from candidates, preferring proximity to label and avoiding claimed regions.

    Scoring: lower is better.
    - Distance to label (if known) is the primary signal.
    - Values to the right of or below their label get a bonus (documents
      typically place values right-of or below labels).
    - Values that overlap the label itself are penalised (that's the label,
      not the value).
    - Overlapping a claimed region adds a large penalty.
    - When multiple label positions are given (label_bboxes), the scoring
      picks the best (label, value) pair.
    """
    if not candidates:
        return None

    CLAIMED_PENALTY = 100_000
    LABEL_OVERLAP_PENALTY = 50_000
    WRONG_DIRECTION_PENALTY = 200  # values above or far-left of label

    all_labels = label_bboxes or ([label_bbox] if label_bbox else [])

    def score_with_label(bb: BoundingBox, lbl: BoundingBox) -> float:
        s = _bbox_distance(bb, lbl)
        # Penalise if the candidate overlaps the label itself (likely IS the label)
        if _bbox_overlaps(bb, lbl, margin=2):
            s += LABEL_OVERLAP_PENALTY
        # Bonus for value being to the right of label (same row) or below
        bb_cx, bb_cy = _bbox_center(bb)
        lbl_cx, lbl_cy = _bbox_center(lbl)
        if bb_cx > lbl_cx or bb_cy > lbl_cy + 10:
            # Good position — slight bonus (reduce score)
            s *= 0.8
        elif bb_cy < lbl_cy - 20 and bb_cx < lbl_cx:
            # Value is above-left of label — unlikely
            s += WRONG_DIRECTION_PENALTY
        return s

    def score(bb: BoundingBox) -> float:
        s = 0.0
        if all_labels:
            # Try each label position and take the best score
            s = min(score_with_label(bb, lbl) for lbl in all_labels)
        # Penalty for overlapping already-claimed regions
        for c in claimed:
            if _bbox_overlaps(bb, c):
                s += CLAIMED_PENALTY
        return s

    scored = sorted(candidates, key=score)
    return scored[0] if scored else None



def generate_bounding_boxes(image: Image.Image, result: "ExtractionResult") -> list[BoundingBox]:
    """Generate bounding boxes by matching extracted field values to OCR word positions.

    Only returns bounding boxes when OCR can reliably match at least 30% of
    extracted fields. For handwritten or complex documents where OCR can't read
    the text (but VLM could), returns an empty list — no boxes is better than
    wrong boxes.
    """
    # Count how many fields actually have values
    field_values = _get_field_values(result)
    numeric_fields = _get_numeric_fields(result)
    total_fields_with_values = sum(1 for v in field_values.values() if v) + len(numeric_fields)

    if total_fields_with_values == 0:
        return []

    # Try OCR-based matching
    words = _get_ocr_word_boxes(image)
    if not words:
        logger.info("No OCR words detected — skipping bounding boxes")
        return []

    bboxes = []
    claimed: list[BoundingBox] = []  # track assigned regions to avoid duplicates

    for field_name, value in field_values.items():
        if not value:
            continue
        # Try alternative representations for dates and currency
        if "date" in field_name:
            search_variants = _date_search_variants(value)
        elif field_name == "currency":
            search_variants = _currency_search_variants(value)
        else:
            search_variants = [value]

        # Gather ALL matches across all variants
        all_matches: list[BoundingBox] = []
        for variant in search_variants:
            all_matches.extend(_find_all_text_bboxes(words, variant))

        # Fallback: for multi-word values with no matches, try progressively
        # shorter sub-phrases, then individual distinctive words
        if not all_matches and len(value.split()) > 1:
            tokens = value.split()
            # Try progressively shorter windows, longest first
            for window in range(len(tokens) - 1, 0, -1):
                for start in range(len(tokens) - window + 1):
                    sub = " ".join(tokens[start:start + window])
                    # For single words, require length >= 5 to be distinctive
                    min_len = 5 if window == 1 else 4
                    if len(sub) >= min_len:
                        sub_matches = _find_all_text_bboxes(words, sub)
                        all_matches.extend(sub_matches)
                if all_matches:
                    break  # found matches at this window size

        if all_matches:
            # Find ALL label positions for context-aware disambiguation
            label_bboxes = _find_label_bboxes(words, field_name)
            label_bbox = label_bboxes[0] if label_bboxes else None
            best = _pick_best_match(all_matches, label_bbox, claimed, label_bboxes=label_bboxes)
            if best:
                best.field = field_name
                bboxes.append(best)
                claimed.append(best)

    for field_name, num_value in numeric_fields.items():
        search_variants = _numeric_search_variants(num_value)
        all_matches = []
        for variant in search_variants:
            all_matches.extend(_find_all_text_bboxes(words, variant))

        if all_matches:
            label_bboxes = _find_label_bboxes(words, field_name)
            label_bbox = label_bboxes[0] if label_bboxes else None
            best = _pick_best_match(all_matches, label_bbox, claimed, label_bboxes=label_bboxes)
            if best:
                best.field = field_name
                bboxes.append(best)
                claimed.append(best)

    # Line items — try to locate each row by its description and/or amount.
    # Uses claimed list so successive line items with the same amount
    # (e.g. two items at $50) each get their own bbox.
    for idx, li in enumerate(result.line_items):
        li_field = f"line_item_{idx}"

        # Find best unclaimed description match
        desc_bbox = None
        if li.description:
            desc_matches = _find_all_text_bboxes(words, li.description)
            desc_bbox = _pick_best_match(desc_matches, None, claimed) if desc_matches else None

        # Find best unclaimed amount match
        amt_bbox = None
        if li.amount is not None and li.amount != 0:
            amt_matches = []
            for variant in _numeric_search_variants(li.amount):
                amt_matches.extend(_find_all_text_bboxes(words, variant))
            if amt_matches:
                # Prefer an amount on the same line as the description
                amt_bbox = _pick_best_match(amt_matches, desc_bbox, claimed)

        if desc_bbox and amt_bbox:
            # Only merge if both are on roughly the same line (within 60px at 200 DPI)
            desc_mid_y = (desc_bbox.y1 + desc_bbox.y2) / 2
            amt_mid_y = (amt_bbox.y1 + amt_bbox.y2) / 2
            if abs(desc_mid_y - amt_mid_y) < 60:
                merged = BoundingBox(
                    field=li_field,
                    x1=min(desc_bbox.x1, amt_bbox.x1),
                    y1=min(desc_bbox.y1, amt_bbox.y1),
                    x2=max(desc_bbox.x2, amt_bbox.x2),
                    y2=max(desc_bbox.y2, amt_bbox.y2),
                )
                bboxes.append(merged)
                claimed.append(merged)
            else:
                desc_bbox.field = li_field
                bboxes.append(desc_bbox)
                claimed.append(desc_bbox)
        elif desc_bbox:
            desc_bbox.field = li_field
            bboxes.append(desc_bbox)
            claimed.append(desc_bbox)
        elif amt_bbox:
            amt_bbox.field = li_field
            bboxes.append(amt_bbox)
            claimed.append(amt_bbox)

    # Only return OCR matches if we matched a meaningful portion of fields.
    # Sparse matches (1-2 out of 10) are usually noise on handwritten docs.
    match_ratio = len(bboxes) / total_fields_with_values if total_fields_with_values > 0 else 0

    # Log which fields were NOT matched for diagnostics
    matched_fields = {bb.field for bb in bboxes}
    all_field_names = set(k for k, v in field_values.items() if v) | set(numeric_fields.keys())
    unmatched = all_field_names - matched_fields
    if unmatched:
        logger.info(f"Unmatched fields: {sorted(unmatched)}")

    if match_ratio < 0.2:
        logger.info(
            f"OCR matched only {len(bboxes)}/{total_fields_with_values} fields "
            f"({match_ratio:.0%}) — too few, skipping bounding boxes"
        )
        return []

    logger.info(f"OCR matched {len(bboxes)}/{total_fields_with_values} fields ({match_ratio:.0%})")
    return bboxes


def _get_field_values(result: "ExtractionResult") -> dict:
    """Get string field name→value mapping for bbox matching.

    Prefers extracted_fields (works for any document type).
    Falls back to legacy named fields for older responses.
    """
    # Primary: extracted_fields dict (dynamic, works for all use cases)
    if result.extracted_fields:
        return {
            k: str(v) for k, v in result.extracted_fields.items()
            if v is not None and not isinstance(v, (int, float))
        }

    # Fallback: legacy named fields (invoice-specific)
    return {
        "vendor_name": result.vendor_name,
        "vendor_address": result.vendor_address,
        "invoice_number": result.invoice_number,
        "invoice_date": result.invoice_date,
        "due_date": result.due_date,
        "currency": result.currency,
        "customer_name": result.customer_name,
        "customer_address": result.customer_address,
        "payment_terms": result.payment_terms,
    }


def _get_numeric_fields(result: "ExtractionResult") -> dict:
    """Get numeric field name→value mapping for bbox matching.

    Prefers extracted_fields (works for any document type).
    Falls back to legacy named fields for older responses.
    """
    # Primary: extracted_fields dict (dynamic, works for all use cases)
    if result.extracted_fields:
        return {
            k: v for k, v in result.extracted_fields.items()
            if isinstance(v, (int, float)) and v is not None
        }

    # Fallback: legacy named fields
    fields = {}
    if result.subtotal is not None:
        fields["subtotal"] = result.subtotal
    if result.tax_amount is not None:
        fields["tax_amount"] = result.tax_amount
    if result.total_amount is not None:
        fields["total_amount"] = result.total_amount
    return fields


def _numeric_search_variants(value: float) -> list[str]:
    """Generate multiple string representations of a number for OCR matching."""
    variants = []
    # Integer form if it's a whole number: "353"
    if value == int(value):
        variants.append(str(int(value)))
    # Standard float: "353.0"
    variants.append(str(value))
    # Two-decimal: "353.00"
    variants.append(f"{value:.2f}")
    # Comma-separated thousands (Western): "1,250.00"
    if value >= 1000:
        variants.append(f"{value:,.2f}")
        variants.append(f"{value:,.0f}" if value == int(value) else f"{value:,.2f}")
    # Indian lakh/crore grouping: "1,23,456.00"
    if value >= 1000:
        int_part = int(value)
        dec_part = f"{value:.2f}".split(".")[1]
        s = str(int_part)
        if len(s) > 3:
            last3 = s[-3:]
            rest = s[:-3]
            # Group remaining digits in pairs from right
            groups = []
            while rest:
                groups.append(rest[-2:] if len(rest) >= 2 else rest)
                rest = rest[:-2]
            groups.reverse()
            indian = ",".join(groups) + "," + last3
        else:
            indian = s
        variants.append(f"{indian}.{dec_part}")
        if value == int(value):
            variants.append(indian)
    return list(dict.fromkeys(variants))  # deduplicate, preserve order


# ---------------------------------------------------------------------------
# Prompt Library — per use-case extraction prompts
# ---------------------------------------------------------------------------

# The library maps (stage, use_case) → prompt text.
# If a use_case-specific prompt doesn't exist, falls back to "_default".

EXTRACTION_PROMPTS: dict[str, str] = {}

EXTRACTION_PROMPTS["invoice"] = """You are a document extraction AI. Extract ALL structured data from this invoice, receipt, or financial document.

Return ONLY valid JSON with this schema (use null for missing fields):
{
  "document_type": "invoice",
  "fields": {
    "vendor_name": "string",
    "vendor_address": "string",
    "invoice_number": "string or null",
    "invoice_date": "string (YYYY-MM-DD)",
    "due_date": "string (YYYY-MM-DD) or null",
    "currency": "string (ISO 4217) or null",
    "subtotal": number or null,
    "tax_amount": number or null,
    "tax_rate": "string (e.g. 18% GST) or null",
    "total_amount": number,
    "payment_terms": "string or null",
    "customer_name": "string or null",
    "customer_address": "string or null"
  },
  "line_items": [{"description": "string", "quantity": number or null, "unit_price": number or null, "amount": number or null}]
}

IMPORTANT RULES:
- Be precise with numbers. Parse dates into YYYY-MM-DD format.
- currency: ONLY set this if a currency symbol or code is VISIBLE. Do NOT infer.
- line_items: Extract ONLY what is explicitly written. Do NOT invent quantity or unit_price. Never default quantity to 1.
- Subtotal: only set if the document explicitly labels a subtotal. Do not copy total_amount into subtotal.
- Include ANY additional visible fields in "fields" even if not listed above."""

EXTRACTION_PROMPTS["tax_form"] = """You are a document extraction AI specializing in tax forms, challans, and government tax documents.

Extract ALL structured data from this tax document.

Return ONLY valid JSON:
{
  "document_type": "tax_form",
  "fields": {
    "form_type": "string (e.g. Form 26QB, Challan 280, Form 16, TDS Certificate)",
    "pan": "string or null",
    "tan": "string or null",
    "assessment_year": "string (e.g. 2026-27) or null",
    "financial_year": "string or null",
    "nature_of_payment": "string or null",
    "section": "string (e.g. 194IA, 234E) or null",
    "deductor_name": "string or null",
    "deductee_name": "string or null",
    "amount_paid": number or null,
    "tds_amount": number or null,
    "surcharge": number or null,
    "education_cess": number or null,
    "total_tax_deposited": number or null,
    "date_of_payment": "string (YYYY-MM-DD) or null",
    "date_of_deduction": "string (YYYY-MM-DD) or null",
    "acknowledgement_number": "string or null",
    "challan_number": "string or null",
    "bsr_code": "string or null",
    "demand_reference_number": "string or null"
  },
  "line_items": []
}

IMPORTANT RULES:
- Include ALL visible fields, even those not listed above — add them to "fields".
- Use snake_case for any extra field names, derived from the label in the document.
- Parse dates into YYYY-MM-DD format.
- Numbers should be numeric types, not strings.
- PAN and TAN should be uppercase strings."""

EXTRACTION_PROMPTS["receipt"] = """You are a document extraction AI. Extract ALL structured data from this receipt or purchase record.

Return ONLY valid JSON:
{
  "document_type": "receipt",
  "fields": {
    "store_name": "string",
    "store_address": "string or null",
    "receipt_number": "string or null",
    "date": "string (YYYY-MM-DD)",
    "currency": "string (ISO 4217) or null",
    "subtotal": number or null,
    "tax_amount": number or null,
    "tax_rate": "string or null",
    "total_amount": number,
    "payment_method": "string or null",
    "card_last_four": "string or null"
  },
  "line_items": [{"description": "string", "quantity": number or null, "unit_price": number or null, "amount": number or null}]
}

IMPORTANT RULES:
- Include ALL visible fields, even those not listed — add them to "fields".
- currency: ONLY set if a symbol or code is VISIBLE. Do NOT infer.
- line_items: Extract ONLY what is explicitly written. Never invent quantity.
- Parse dates into YYYY-MM-DD format."""

EXTRACTION_PROMPTS["_default"] = """You are a document extraction AI. Analyze this document and extract ALL visible structured data.

Steps:
1. Identify the document type (e.g. invoice, receipt, tax_form, contract, bank_statement, id_card, letter, application, certificate, etc.)
2. Extract every labeled field visible in the document as key-value pairs.
3. Extract any tabular data as line_items.

Return ONLY valid JSON:
{
  "document_type": "string",
  "fields": {
    "field_name": "value",
    ...
  },
  "line_items": [
    {"column_name": "value", ...},
    ...
  ]
}

IMPORTANT RULES:
- Use snake_case for field names, derived from the label text in the document.
- Parse dates into YYYY-MM-DD format.
- Numbers should be numeric types, not strings.
- Include ALL visible fields — headers, form fields, metadata, identifiers, dates, amounts.
- For tables, use column headers as keys in each row object.
- Do NOT infer or guess values. Only extract what is explicitly visible."""

# ---------------------------------------------------------------------------
# Validation Prompts
# ---------------------------------------------------------------------------

VALIDATION_PROMPTS: dict[str, str] = {}

VALIDATION_PROMPTS["invoice"] = """You are a document validation AI. You will be given:
1. An image of the original document.
2. A JSON object of extracted data.

Verify EVERY extracted field against the document image. For each field, check:
- Is the value present in the document?
- Is it accurately transcribed (exact text, numbers, dates)?
- Are numeric fields consistent (subtotal + tax = total, qty × price = amount)?

Return ONLY valid JSON:
{
  "checks": [
    {
      "field": "field_name",
      "status": "pass" | "fail" | "warn",
      "message": "short explanation",
      "expected": "what the document shows (if different)",
      "actual": "what was extracted"
    }
  ],
  "summary": "short overall assessment"
}

Focus on: amounts, dates, vendor/customer names, invoice numbers, tax calculations.
Flag any field where the extracted value doesn't match what's visible in the image."""

VALIDATION_PROMPTS["tax_form"] = """You are a document validation AI specializing in tax forms and government documents.

You will be given:
1. An image of the original tax document.
2. A JSON object of extracted data.

Verify EVERY extracted field against the document image. For tax forms, pay special attention to:
- PAN/TAN format validity (10 chars: 5 letters, 4 digits, 1 letter)
- Assessment year / financial year consistency
- TDS amount + surcharge + cess = total tax deposited
- Date fields match what's printed
- Section numbers and challan/acknowledgement numbers are exact

Return ONLY valid JSON:
{
  "checks": [
    {
      "field": "field_name",
      "status": "pass" | "fail" | "warn",
      "message": "short explanation",
      "expected": "what the document shows (if different)",
      "actual": "what was extracted"
    }
  ],
  "summary": "short overall assessment"
}"""

VALIDATION_PROMPTS["receipt"] = """You are a document validation AI. You will be given:
1. An image of the original receipt.
2. A JSON object of extracted data.

Verify EVERY extracted field against the receipt image. For receipts, check:
- Store name and details match the header
- Line item amounts are correct
- Subtotal/tax/total arithmetic
- Date and receipt number accuracy
- Payment method matches what's printed

Return ONLY valid JSON:
{
  "checks": [
    {
      "field": "field_name",
      "status": "pass" | "fail" | "warn",
      "message": "short explanation",
      "expected": "what the document shows (if different)",
      "actual": "what was extracted"
    }
  ],
  "summary": "short overall assessment"
}"""

VALIDATION_PROMPTS["_default"] = """You are a document validation AI. You will be given:
1. An image of the original document.
2. A JSON object of extracted data.

Verify EVERY extracted field against the document image:
- Is each field value present and accurately transcribed?
- Are numeric fields consistent with each other?
- Are dates in the correct format and matching the document?
- Are identifiers (numbers, codes, references) exact?

Return ONLY valid JSON:
{
  "checks": [
    {
      "field": "field_name",
      "status": "pass" | "fail" | "warn",
      "message": "short explanation",
      "expected": "what the document shows (if different)",
      "actual": "what was extracted"
    }
  ],
  "summary": "short overall assessment"
}"""


# Stage-level prompt library: maps (stage, use_case) → prompt.
PROMPT_LIBRARY: dict[str, dict[str, str]] = {
    "extraction": EXTRACTION_PROMPTS,
    "validation": VALIDATION_PROMPTS,
}


def get_prompt(stage: str, use_case: str) -> str:
    """Look up the prompt for a (stage, use_case) pair.

    Priority:
      1. Active prompt from Supabase prompt_versions table (if configured)
      2. Hardcoded PROMPT_LIBRARY dict (in-code defaults)
      3. _default prompt for the stage
    """
    # Try DB first — gracefully returns None when Supabase is not configured
    try:
        from db.supabase import get_active_prompt
        db_prompt = get_active_prompt(use_case, stage)
        if db_prompt:
            return db_prompt
    except Exception:
        pass  # Fall through to hardcoded prompts

    stage_prompts = PROMPT_LIBRARY.get(stage, {})
    return stage_prompts.get(use_case, stage_prompts.get("_default", ""))


# Legacy alias kept so callers that reference EXTRACTION_PROMPT still work
EXTRACTION_PROMPT = EXTRACTION_PROMPTS["invoice"]


def image_to_base64(image: Image.Image, max_size: int = 1024) -> str:
    """Convert PIL Image to base64 string, resizing if needed.

    IMPORTANT: Works on a copy so the original image is never mutated.
    This preserves full resolution for downstream OCR bounding-box generation.
    """
    img = image.copy()
    if max(img.size) > max_size:
        img.thumbnail((max_size, max_size), Image.LANCZOS)
    if img.mode == "RGBA":
        img = img.convert("RGB")
    buffer = BytesIO()
    img.save(buffer, format="JPEG", quality=85)
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def _parse_vlm_json(raw: str) -> dict:
    """Extract JSON from a VLM response, handling markdown fences."""
    json_match = re.search(r"```(?:json)?\s*([\s\S]*?)```", raw)
    if json_match:
        return json.loads(json_match.group(1).strip())
    json_match = re.search(r"\{[\s\S]*\}", raw)
    if json_match:
        return json.loads(json_match.group(0))
    return json.loads(raw)


def _vlm_data_to_result(data: dict, model_id: str, use_case: str) -> ExtractionResult:
    """Convert parsed VLM JSON into an ExtractionResult.

    Handles two response shapes:
      1. New format: {"document_type": "...", "fields": {...}, "line_items": [...]}
      2. Legacy flat format: {"vendor_name": "...", "total_amount": 100, ...}
    """
    short_model = model_id.split("/")[-1]

    # Detect response shape
    fields = data.get("fields", {})
    doc_type = data.get("document_type", use_case)

    if not fields:
        # Legacy flat format — all keys are fields (except line_items)
        fields = {k: v for k, v in data.items()
                  if k not in ("line_items", "document_type", "fields")}

    result = ExtractionResult(
        extraction_method=f"vlm:{short_model}",
        confidence=0.92,
        document_type=doc_type,
        extracted_fields=fields,
    )

    # Populate legacy named fields from fields dict (backward compat for invoice)
    result.vendor_name = str(fields.get("vendor_name", "") or "")
    result.vendor_address = str(fields.get("vendor_address", "") or "")
    result.invoice_number = str(fields.get("invoice_number", "") or "")
    result.invoice_date = str(fields.get("invoice_date", "") or "")
    result.due_date = str(fields.get("due_date", "") or "")
    result.currency = str(fields.get("currency", "") or "")
    result.subtotal = fields.get("subtotal")
    result.tax_amount = fields.get("tax_amount")
    result.tax_rate = str(fields.get("tax_rate", "") or "")
    result.total_amount = fields.get("total_amount")
    result.payment_terms = str(fields.get("payment_terms", "") or "")
    result.customer_name = str(fields.get("customer_name", "") or "")
    result.customer_address = str(fields.get("customer_address", "") or "")

    # Parse line items
    for item in data.get("line_items", []):
        if isinstance(item, dict):
            result.line_items.append(LineItem(
                description=item.get("description", ""),
                quantity=item.get("quantity"),
                unit_price=item.get("unit_price"),
                amount=item.get("amount"),
            ))

    return result


def extract_with_vlm(
    image: Image.Image,
    hf_token: str,
    use_case: str = "invoice",
) -> Optional[ExtractionResult]:
    """Extract document data using a Vision-Language Model via HF Inference API.

    Selects the prompt from the prompt library based on use_case.
    Falls back to the generic prompt when no use_case-specific prompt exists.
    """
    prompt = get_prompt("extraction", use_case)
    if not prompt:
        logger.warning(f"No extraction prompt for use_case={use_case}, using default")
        prompt = get_prompt("extraction", "_default")

    # Try multiple models in order of preference
    models = [
        "Qwen/Qwen2.5-VL-72B-Instruct",
        "Qwen/Qwen2.5-VL-7B-Instruct",
        "Qwen/Qwen2.5-VL-3B-Instruct",
    ]

    img_b64 = image_to_base64(image)
    data_url = f"data:image/jpeg;base64,{img_b64}"
    errors = []

    for model_id in models:
        try:
            client = InferenceClient(api_key=hf_token)

            response = client.chat_completion(
                model=model_id,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {
                                "type": "image_url",
                                "image_url": {"url": data_url},
                            },
                        ],
                    }
                ],
                max_tokens=2048,
                temperature=0.1,
            )

            result_text = response.choices[0].message.content
            data = _parse_vlm_json(result_text)
            return _vlm_data_to_result(data, model_id, use_case)

        except Exception as e:
            short_model = model_id.split("/")[-1]
            error_msg = str(e)[:150]
            errors.append(f"{short_model}: {type(e).__name__}({error_msg})")
            logger.warning(f"VLM extraction failed with {model_id}: {e}")
            continue

    return ExtractionResult(
        extraction_method="failed:vlm_all_models",
        confidence=0.0,
        raw_text="; ".join(errors) if errors else "all models failed silently",
    )


# ---------------------------------------------------------------------------
# OCR + Regex Extraction (Fallback) - CPU only, zero cost
# ---------------------------------------------------------------------------

def extract_text_ocr(image: Image.Image) -> str:
    """Extract text from image using Tesseract OCR."""
    try:
        import pytesseract
        text = pytesseract.image_to_string(image, config="--psm 3")
        return text
    except ImportError:
        logger.warning("pytesseract not available")
        return ""
    except Exception as e:
        logger.warning(f"OCR failed: {e}")
        return ""


def parse_amount(text: str) -> Optional[float]:
    """Parse a monetary amount from text."""
    text = text.replace(",", "").replace(" ", "")
    match = re.search(r"[\d]+\.?\d*", text)
    if match:
        try:
            return float(match.group())
        except ValueError:
            pass
    return None


def detect_currency(text: str) -> str:
    """Detect currency from text symbols and keywords."""
    currency_patterns = {
        "INR": [r"₹", r"Rs\.?", r"INR", r"rupee"],
        "USD": [r"\$", r"USD", r"dollar"],
        "EUR": [r"€", r"EUR", r"euro"],
        "GBP": [r"£", r"GBP", r"pound"],
    }
    for currency, patterns in currency_patterns.items():
        for p in patterns:
            if re.search(p, text, re.IGNORECASE):
                return currency
    return "USD"  # default


def extract_with_regex(text: str) -> ExtractionResult:
    """Extract structured data from OCR text using regex patterns."""
    result = ExtractionResult(
        raw_text=text,
        extraction_method="ocr:tesseract+regex",
        confidence=0.65,
    )

    # Invoice number — require a separator (No, #, :) between label and value
    inv_patterns = [
        r"(?:invoice|inv|bill)\s*(?:#|no\.?|number)\s*:?\s*([A-Z0-9\-/]+)",
        r"(?:invoice|inv)\s*:\s*([A-Z0-9\-/]+)",
        r"(?:#|no\.?)\s*:?\s*([A-Z0-9\-/]+)",
    ]
    for p in inv_patterns:
        m = re.search(p, text, re.IGNORECASE)
        if m:
            result.invoice_number = m.group(1).strip()
            break

    # Dates
    date_patterns = [
        r"(\d{1,2}[/\-\.]\d{1,2}[/\-\.]\d{2,4})",
        r"(\d{4}[/\-\.]\d{1,2}[/\-\.]\d{1,2})",
        r"(\w+\s+\d{1,2},?\s+\d{4})",
    ]
    dates_found = []
    for p in date_patterns:
        dates_found.extend(re.findall(p, text))
    if dates_found:
        result.invoice_date = dates_found[0]
        if len(dates_found) > 1:
            result.due_date = dates_found[1]

    # Total amount — try grand total / total due first, then plain "total"
    # but skip lines that say "subtotal" or specific tax labels
    total_patterns = [
        r"(?:grand\s*total|amount\s*due|balance\s*due|total\s*amount)\s*:?\s*(?:INR|USD|EUR|GBP|[₹$€£])?\s*([\d,]+\.?\d*)",
        r"(?<!sub)total\s*:?\s*(?:INR|USD|EUR|GBP|[₹$€£])?\s*([\d,]+\.?\d*)",
    ]
    for p in total_patterns:
        matches = list(re.finditer(p, text, re.IGNORECASE | re.MULTILINE))
        if matches:
            # Take the last match (usually the grand total at the bottom)
            result.total_amount = parse_amount(matches[-1].group(1))
            break

    # Subtotal
    m = re.search(r"(?:subtotal|sub\s*total)\s*:?\s*[₹$€£]?\s*([\d,]+\.?\d*)", text, re.IGNORECASE)
    if m:
        result.subtotal = parse_amount(m.group(1))

    # Tax
    tax_patterns = [
        r"(?:tax|gst|vat|igst|cgst|sgst)\s*(?:\(?\d+%?\)?)?\s*:?\s*[₹$€£]?\s*([\d,]+\.?\d*)",
    ]
    for p in tax_patterns:
        m = re.search(p, text, re.IGNORECASE)
        if m:
            result.tax_amount = parse_amount(m.group(1))
            break

    # Tax rate — also match "CGST (9%)" and "GST @18%" patterns
    tax_rate_patterns = [
        r"(\d+(?:\.\d+)?)\s*%\s*(?:tax|gst|vat|igst)",
        r"(?:tax|gst|vat|cgst|sgst|igst)\s*[@(]?\s*(\d+(?:\.\d+)?)\s*%",
    ]
    for p in tax_rate_patterns:
        m = re.search(p, text, re.IGNORECASE)
        if m:
            result.tax_rate = f"{m.group(1)}%"
            break

    # Currency
    result.currency = detect_currency(text)

    # Vendor name (heuristic: first line or line before address)
    lines = [l.strip() for l in text.split("\n") if l.strip()]
    if lines:
        result.vendor_name = lines[0]

    return result


# ---------------------------------------------------------------------------
# Main Extraction Pipeline
# ---------------------------------------------------------------------------

def extract_invoice(
    image: Image.Image,
    hf_token: Optional[str] = None,
    force_ocr: bool = False,
    use_case: str = "invoice",
) -> ExtractionResult:
    """
    Main extraction pipeline. Tries VLM first, falls back to OCR+regex.

    Args:
        image: PIL Image of the document
        hf_token: Hugging Face API token (for VLM extraction)
        force_ocr: Skip VLM and use OCR only
        use_case: Document type for prompt selection (invoice, tax_form, receipt, etc.)

    Returns:
        ExtractionResult with all extracted fields
    """
    start = datetime.now()

    result = None
    vlm_errors = ""

    # Try VLM extraction first (if token available and not forced to OCR)
    if hf_token and not force_ocr:
        vlm_result = extract_with_vlm(image, hf_token, use_case=use_case)
        if vlm_result and not vlm_result.extraction_method.startswith("failed"):
            result = vlm_result
        else:
            # VLM failed — capture diagnostics, fall through to OCR
            vlm_errors = vlm_result.raw_text if vlm_result else "unknown"
            logger.warning(f"VLM failed ({vlm_errors}), trying OCR fallback")
    elif not hf_token and not force_ocr:
        vlm_errors = "no_token"
        logger.warning("HF_TOKEN not set — skipping VLM extraction, using OCR fallback only")

    # Fallback to OCR + regex
    if result is None:
        text = extract_text_ocr(image)
        if text.strip():
            result = extract_with_regex(text)
            if vlm_errors:
                # Note that this was a fallback result
                result.extraction_method += f" (vlm failed: {vlm_errors})"
        else:
            method = f"failed:vlm({vlm_errors})+ocr_empty" if vlm_errors else "failed:ocr_empty"
            result = ExtractionResult(
                extraction_method=method,
                confidence=0.0,
            )

    # Generate bounding boxes by matching extracted values to OCR word positions.
    # Falls back to layout-based estimation for handwritten documents.
    try:
        result.bounding_boxes = generate_bounding_boxes(image, result)
    except Exception as e:
        logger.warning(f"Bounding box generation failed: {e}")

    elapsed = (datetime.now() - start).total_seconds() * 1000
    result.processing_time_ms = int(elapsed)

    return result


def extract_from_pdf(
    pdf_bytes: bytes,
    hf_token: Optional[str] = None,
    force_ocr: bool = False,
    use_case: str = "invoice",
) -> list[ExtractionResult]:
    """Extract from PDF - converts each page to image and extracts.

    The first page's rendered image is attached as ``page_image`` (a
    data-URI) so the frontend can display the PDF and overlay bounding
    boxes without needing a separate PDF renderer.
    """
    try:
        import fitz  # PyMuPDF
    except ImportError:
        raise ImportError("PyMuPDF (fitz) is required for PDF processing. Install with: pip install PyMuPDF")

    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    results = []

    for page_num in range(min(len(doc), 10)):  # Max 10 pages
        page = doc[page_num]
        pix = page.get_pixmap(dpi=200)
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        result = extract_invoice(img, hf_token=hf_token, force_ocr=force_ocr, use_case=use_case)

        # Attach rendered page image so the viewer can display and overlay bboxes
        buf = BytesIO()
        img.save(buf, format="PNG", optimize=True)
        b64 = base64.b64encode(buf.getvalue()).decode()
        result.page_image = f"data:image/png;base64,{b64}"

        results.append(result)

    doc.close()
    return results


# ---------------------------------------------------------------------------
# VLM-based Semantic Validation
# ---------------------------------------------------------------------------

def validate_with_vlm(
    image: Image.Image,
    extraction_data: dict,
    hf_token: str,
    model_id: str = "Qwen/Qwen2.5-VL-72B-Instruct",
    use_case: str = "_default",
) -> list[dict]:
    """Use a VLM to cross-check extracted data against the source image.

    Returns a list of check dicts: {field, status, message, expected, actual}.
    """
    if not hf_token:
        return [{
            "name": "VLM validation",
            "status": "skipped",
            "message": "HF_TOKEN not set — cannot run VLM validation",
        }]

    prompt = get_prompt("validation", use_case)
    if not prompt:
        prompt = get_prompt("validation", "_default")

    # Build the extraction summary for the prompt
    fields = extraction_data.get("extracted_fields", {})
    if not fields:
        # Legacy format
        fields = {k: v for k, v in extraction_data.items()
                  if k not in ("line_items", "bounding_boxes", "raw_text",
                               "extraction_method", "confidence", "processing_time_ms",
                               "page_image", "page_images", "run_id", "error",
                               "document_type", "extracted_fields")}

    extraction_json = json.dumps({
        "fields": fields,
        "line_items": extraction_data.get("line_items", []),
    }, indent=2)

    full_prompt = f"{prompt}\n\nExtracted data to validate:\n```json\n{extraction_json}\n```"

    try:
        b64 = image_to_base64(image)
        client = InferenceClient(provider="auto", api_key=hf_token)

        response = client.chat_completion(
            model=model_id,
            messages=[{
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
                    {"type": "text", "text": full_prompt},
                ],
            }],
            max_tokens=2048,
        )

        raw = response.choices[0].message.content
        data = _parse_vlm_json(raw)
        checks = data.get("checks", [])

        # Normalize check format
        result = []
        for c in checks:
            result.append({
                "name": f"VLM: {c.get('field', 'unknown')}",
                "status": c.get("status", "warn"),
                "message": c.get("message", ""),
                "expected": c.get("expected"),
                "actual": c.get("actual"),
            })

        return result if result else [{
            "name": "VLM validation",
            "status": "pass",
            "message": data.get("summary", "All fields appear correct"),
        }]

    except Exception as e:
        logger.warning(f"VLM validation failed: {e}")
        return [{
            "name": "VLM validation",
            "status": "skipped",
            "message": f"VLM validation unavailable: {str(e)[:100]}",
        }]


def run_dynamic_math_checks(data: dict) -> list[dict]:
    """Run math/consistency checks dynamically based on extracted_fields.

    Works for any document type by detecting numeric relationships.
    """
    checks = []
    fields = data.get("extracted_fields", {})
    if not fields:
        # Fall back to top-level keys
        fields = data

    line_items = data.get("line_items", [])

    # ── Collect numeric fields ──
    numeric = {}
    for k, v in fields.items():
        if v is not None:
            try:
                numeric[k] = float(v)
            except (ValueError, TypeError):
                pass

    # ── Rule 1: subtotal + tax-like = total-like ──
    subtotal_key = next((k for k in numeric if "subtotal" in k), None)
    total_key = next((k for k in numeric if "total" in k and "sub" not in k), None)
    tax_keys = [k for k in numeric if any(t in k for t in ("tax", "cess", "surcharge", "duty", "vat", "gst"))]

    if subtotal_key and total_key and tax_keys:
        tax_sum = sum(numeric[k] for k in tax_keys)
        expected = numeric[subtotal_key] + tax_sum
        diff = abs(expected - numeric[total_key])
        checks.append({
            "name": "Subtotal + taxes = Total",
            "status": "pass" if diff < 0.02 else "warn" if diff < 1.0 else "fail",
            "message": f"{subtotal_key} + {' + '.join(tax_keys)} = {total_key}",
            "expected": numeric[total_key],
            "actual": round(expected, 2),
        })

    # ── Rule 2: Line item amounts sum to subtotal or total ──
    if line_items:
        item_sum = sum(
            float(item.get("amount", 0) or 0) for item in line_items
        )
        compare_key = subtotal_key or total_key
        if compare_key and item_sum > 0:
            diff = abs(item_sum - numeric[compare_key])
            checks.append({
                "name": "Line items sum",
                "status": "pass" if diff < 0.02 else "warn" if diff < 1.0 else "fail",
                "message": f"Sum of line item amounts vs {compare_key}",
                "expected": numeric[compare_key],
                "actual": round(item_sum, 2),
            })

    # ── Rule 3: Each line item qty × unit_price = amount ──
    for i, item in enumerate(line_items):
        qty = item.get("quantity")
        price = item.get("unit_price")
        amount = item.get("amount")
        if qty is not None and price is not None and amount is not None:
            try:
                expected = round(float(qty) * float(price), 2)
                diff = abs(expected - float(amount))
                checks.append({
                    "name": f"Line {i+1} math",
                    "status": "pass" if diff < 0.02 else "fail",
                    "message": f"qty ({qty}) × unit_price ({price}) = amount",
                    "expected": float(amount),
                    "actual": expected,
                })
            except (ValueError, TypeError):
                pass

    # ── Rule 4: Total positive ──
    if total_key:
        checks.append({
            "name": "Total > 0",
            "status": "pass" if numeric[total_key] > 0 else "fail",
            "message": f"{total_key} should be positive",
            "expected": "> 0",
            "actual": numeric[total_key],
        })

    # ── Rule 5: PAN/TAN format (tax forms) ──
    for k, v in fields.items():
        if isinstance(v, str):
            if k == "pan" and v:
                valid = bool(re.match(r'^[A-Z]{5}[0-9]{4}[A-Z]$', v))
                checks.append({
                    "name": "PAN format",
                    "status": "pass" if valid else "fail",
                    "message": "PAN should be 5 letters + 4 digits + 1 letter",
                    "expected": "XXXXX9999X",
                    "actual": v,
                })
            elif k == "tan" and v:
                valid = bool(re.match(r'^[A-Z]{4}[0-9]{5}[A-Z]$', v))
                checks.append({
                    "name": "TAN format",
                    "status": "pass" if valid else "fail",
                    "message": "TAN should be 4 letters + 5 digits + 1 letter",
                    "expected": "XXXX99999X",
                    "actual": v,
                })

    if not checks:
        checks.append({
            "name": "Insufficient data",
            "status": "skipped",
            "message": "Not enough numeric or validatable fields to run checks",
        })

    return checks