#!/usr/bin/env python3
"""Local-first receipt/document importer for knowledge/receipts.

Pipeline:
1. Keep the original image/PDF immutable in knowledge/receipts/images/YY-MM-images/.
2. Generate local OCR derivatives in knowledge/receipts/ocr/<sha-prefix>/.
3. OCR locally with Tesseract/OCRmyPDF/Poppler.
4. Parse likely receipt fields with conservative heuristics.
5. Insert a searchable row into receipts.sqlite, or update an existing receipt id.

No cloud APIs. No deletion. Designed as a wrapper Soren can call quickly from chat.
"""

from __future__ import annotations

import argparse
import datetime as dt
import difflib
import hashlib
import json
import os
import re
import shutil
import sqlite3
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any

try:
    from PIL import Image, ImageEnhance, ImageOps
except Exception as e:  # pragma: no cover
    print(f"Pillow/PIL required: {e}", file=sys.stderr)
    raise SystemExit(2)

ROOT = Path(__file__).resolve().parents[1]
DB = ROOT / "receipts.sqlite"
IMAGES = ROOT / "images"
OCR_ROOT = ROOT / "ocr"
TEXT_EXTS = {".txt"}
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".tif", ".tiff", ".bmp"}
PDF_EXTS = {".pdf"}


def run(cmd: list[str], *, capture: bool = False, check: bool = True, timeout: int = 120) -> subprocess.CompletedProcess[str]:
    try:
        return subprocess.run(cmd, text=True, capture_output=capture, check=check, timeout=timeout)
    except subprocess.CalledProcessError as e:
        if e.stdout:
            print(e.stdout, end="")
        if e.stderr:
            print(e.stderr, end="", file=sys.stderr)
        raise SystemExit(e.returncode)


def have(name: str) -> bool:
    return shutil.which(name) is not None


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def now_local() -> dt.datetime:
    return dt.datetime.now().astimezone()


def month_dir(ts: dt.datetime) -> Path:
    return IMAGES / f"{ts:%y-%m}-images"


def slug(s: str) -> str:
    s = re.sub(r"[^A-Za-z0-9._-]+", "-", s.strip()).strip("-_.")
    return s[:80] or "receipt"


def copy_original(src: Path, *, category: str, captured: dt.datetime, sha: str) -> Path:
    ext = src.suffix.lower() or ".jpg"
    dest_dir = month_dir(captured)
    dest_dir.mkdir(parents=True, exist_ok=True)
    cat = slug(category.lower())
    dest = dest_dir / f"{captured:%Y-%m-%d_%H%M}_{cat}_{sha[:8]}{ext}"
    if not dest.exists():
        shutil.copy2(src, dest)
    return dest


def image_derivatives(src: Path, ocr_dir: Path) -> list[Path]:
    ocr_dir.mkdir(parents=True, exist_ok=True)
    im = Image.open(src).convert("RGB")
    outputs: list[Path] = []

    # 1. Original as PNG, useful for consistent tesseract input.
    original = ocr_dir / "original.png"
    im.save(original)
    outputs.append(original)

    # 2. Grayscale + contrast + upscale. Good for screenshots and receipts.
    g = ImageOps.grayscale(im)
    g = ImageEnhance.Contrast(g).enhance(2.2)
    up = g.resize((g.width * 2, g.height * 2))
    gray = ocr_dir / "gray_contrast_2x.png"
    up.save(gray)
    outputs.append(gray)

    # 3. Binary-ish threshold, helps faint receipt paper.
    bw = up.point(lambda p: 255 if p > 175 else 0)
    bw_path = ocr_dir / "threshold_2x.png"
    bw.save(bw_path)
    outputs.append(bw_path)

    # 4. Vertical crops, useful for long phone screenshots/receipts.
    if im.height > im.width * 1.2:
        for i in range(4):
            y0 = i * im.height // 4
            y1 = (i + 1) * im.height // 4
            crop = im.crop((0, y0, im.width, y1))
            crop = ImageEnhance.Contrast(ImageOps.grayscale(crop)).enhance(2.0)
            crop = crop.resize((crop.width * 2, crop.height * 2))
            p = ocr_dir / f"crop_{i+1}_2x.png"
            crop.save(p)
            outputs.append(p)

    return outputs


def tesseract_text(image: Path, *, psm: int) -> str:
    if not have("tesseract"):
        raise RuntimeError("tesseract not installed")
    cp = run(["tesseract", str(image), "stdout", "-l", "eng", "--psm", str(psm)], capture=True, check=False)
    return cp.stdout or ""


def ocr_image(src: Path, ocr_dir: Path) -> dict[str, Any]:
    variants = image_derivatives(src, ocr_dir)
    runs = []
    combined_parts: list[str] = []
    for img in variants:
        for psm in (6, 11):
            text = tesseract_text(img, psm=psm)
            cleaned = clean_text(text)
            score = score_text(cleaned)
            runs.append({"file": str(img), "psm": psm, "score": score, "chars": len(cleaned)})
            if cleaned:
                combined_parts.append(f"--- OCR {img.name} psm={psm} ---\n{cleaned}")
    combined = dedupe_lines("\n".join(combined_parts))
    (ocr_dir / "ocr_raw.txt").write_text(combined)
    (ocr_dir / "ocr_runs.json").write_text(json.dumps(runs, indent=2))
    return {"text": combined, "runs": runs, "dir": str(ocr_dir)}


def ocr_pdf(src: Path, ocr_dir: Path) -> dict[str, Any]:
    ocr_dir.mkdir(parents=True, exist_ok=True)
    raw_text = ""
    if have("pdftotext"):
        cp = run(["pdftotext", "-layout", str(src), "-"], capture=True, check=False)
        raw_text = clean_text(cp.stdout or "")
    ocr_pdf_path = None
    if len(raw_text.strip()) < 40 and have("ocrmypdf"):
        ocr_pdf_path = ocr_dir / "ocr.pdf"
        run(["ocrmypdf", "--skip-text", "--deskew", "--rotate-pages", str(src), str(ocr_pdf_path)], check=False, timeout=300)
        if ocr_pdf_path.exists() and have("pdftotext"):
            cp = run(["pdftotext", "-layout", str(ocr_pdf_path), "-"], capture=True, check=False)
            raw_text = clean_text(cp.stdout or raw_text)
    (ocr_dir / "ocr_raw.txt").write_text(raw_text)
    return {"text": raw_text, "runs": [{"pdf_ocr": bool(ocr_pdf_path and ocr_pdf_path.exists())}], "dir": str(ocr_dir)}


def clean_text(text: str) -> str:
    text = text.replace("\r", "\n")
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def dedupe_lines(text: str) -> str:
    seen = set()
    out = []
    for line in text.splitlines():
        key = re.sub(r"\s+", " ", line.strip()).lower()
        if not key:
            if out and out[-1] != "":
                out.append("")
            continue
        # Preserve OCR section headers even when similar.
        if key.startswith("--- ocr") or key not in seen:
            out.append(line.rstrip())
            seen.add(key)
    return "\n".join(out).strip()


def score_text(text: str) -> int:
    score = len(text)
    for word in ("total", "subtotal", "gst", "receipt", "invoice", "order", "payment", "$"):
        score += text.lower().count(word) * 50
    return score


def money_to_float(s: str | None, *, require_dollars: bool = False, last: bool = False) -> float | None:
    if not s:
        return None
    s = s.replace(",", "")
    pat = r"-?\$\s*([0-9]+(?:\.[0-9]{2})?)" if require_dollars else r"-?\$?\s*([0-9]+(?:\.[0-9]{2})?)"
    matches = list(re.finditer(pat, s))
    if not matches:
        return None
    m = matches[-1] if last else matches[0]
    return float(m.group(1))


def parse_date(text: str) -> str | None:
    patterns = [
        (r"\b(20\d{2})[-/](\d{1,2})[-/](\d{1,2})\b", "%Y-%m-%d"),
        (r"\b(\d{1,2})[-/](\d{1,2})[-/](20\d{2})\b", "%d-%m-%Y"),
        (r"\b([A-Z][a-z]+)\s+(\d{1,2}),\s*(20\d{2})\b", "%B %d %Y"),
        (r"\b(\d{1,2})\s+([A-Z][a-z]+)\s+(20\d{2})\b", "%d %B %Y"),
    ]
    for pat, fmt in patterns:
        for m in re.finditer(pat, text):
            val = " ".join(m.groups()) if "%B" in fmt else "-".join(m.groups())
            try:
                return dt.datetime.strptime(val, fmt).date().isoformat()
            except ValueError:
                continue
    return None


def likely_merchant(lines: list[str], text: str = "") -> str | None:
    # Specific useful hint: many web receipts contain the domain even when the logo OCR is noisy.
    domain = re.search(r"\b(?:www\.)?([a-z0-9-]+)\.com\.au\b", text, re.I)
    if domain:
        known = {"nopong": "No Pong"}
        key = domain.group(1).lower().replace("-", "")
        if key in known:
            return known[key]
        return domain.group(1).replace("-", " ").title()

    bad = re.compile(r"^(--- ocr|receipt|tax invoice|invoice|order|checkout|date|total|subtotal|payment|product|qty|amount|thank you)", re.I)
    for line in lines[:30]:
        l = line.strip(" :-|•")
        if len(l) < 3 or bad.search(l) or re.search(r"@|https?://|www\.|\$|\d{4,}", l):
            continue
        # Prefer line with letters, not mostly symbols.
        if sum(ch.isalpha() for ch in l) >= 3:
            return l[:120]
    return None


def parse_fields(text: str) -> dict[str, Any]:
    lines = [l.strip() for l in text.splitlines() if l.strip()]
    lower = text.lower()
    fields: dict[str, Any] = {
        "merchant_name": likely_merchant(lines, text),
        "receipt_date": parse_date(text),
        "currency": "AUD" if "$" in text else None,
    }

    # Order / invoice / transaction refs. Prefer same-line matches, then label-next-line matches.
    ref_patterns = {
        "docket_number": r"(?:order\s*(?:number|no\.?|#)|receipt\s*(?:number|no\.?|#)|docket\s*(?:number|no\.?|#))\s*[:#]?\s*([A-Z0-9-]{4,})",
        "invoice_number": r"invoice\s*(?:number|no\.?|#)?\s*[:#]?\s*([A-Z0-9-]{4,})",
        "transaction_reference": r"(?:transaction|txn|reference|ref)\s*(?:number|no\.?|#)?\s*[:#]?\s*([A-Z0-9-]{4,})",
    }
    for key, pat in ref_patterns.items():
        m = re.search(pat, text, re.I)
        if m and m.group(1).lower() not in {"date", "total", "email", "payment"}:
            fields[key] = m.group(1)
    for i, line in enumerate(lines[:-1]):
        if "order number" in line.lower() and not fields.get("docket_number"):
            # Common receipt layout: ORDER NUMBER: DATE: / 3465186 May 14, 2026
            nxt = lines[i + 1]
            m = re.search(r"\b([0-9]{4,}|[A-Z0-9-]{6,})\b", nxt)
            if m:
                fields["docket_number"] = m.group(1)

    pm = re.search(r"payment\s*method\s*:?\s*([A-Za-z0-9 /_-]{3,40})", text, re.I)
    if pm:
        fields["payment_method"] = pm.group(1).strip()
    else:
        for method in ("PayPal", "Visa", "Mastercard", "EFTPOS", "Cash", "Amex", "Afterpay", "Apple Pay", "Google Pay"):
            if method.lower() in lower:
                fields["payment_method"] = method
                break

    # Amounts. Line-based parsing avoids confusing PRODUCT TOTAL headers and 85G sizes.
    for line in lines:
        low = line.lower()
        amount_last = money_to_float(line, require_dollars=True, last=True)
        amount_first = money_to_float(line, require_dollars=True, last=False)
        if amount_last is None:
            continue
        if "subtotal" in low:
            fields["subtotal"] = amount_last
        if re.search(r"\bgst\b|\btax\b", low):
            fields["gst"] = amount_last
        if re.search(r"^(grand\s+total|amount\s+paid|total)\b", low):
            # For "Total: $99.98 (includes $9.09 GST)", the total is the first amount.
            fields["total"] = amount_first
    # Embedded form: includes $9.09 GST
    gst_m = re.search(r"includes\s+\$\s*([0-9,]+\.[0-9]{2})\s+GST", text, re.I)
    if gst_m:
        fields["gst"] = float(gst_m.group(1).replace(",", ""))

    # Fallback total = largest money amount, conservative.
    if not fields.get("total"):
        monies = [money_to_float(m.group(0), require_dollars=True) for m in re.finditer(r"\$\s*[0-9,]+\.[0-9]{2}", text)]
        monies = [m for m in monies if m is not None]
        if monies:
            fields["total"] = max(monies)

    # Simple item extraction: lines with money, excluding summary/payment lines.
    item_lines = []
    seen_items = set()
    skip = re.compile(r"subtotal|total|gst|tax|shipping|payment|points|redeemed|balance|change|cash|visa|mastercard|email|@", re.I)
    for line in lines:
        if "$" not in line or skip.search(line):
            continue
        amount = money_to_float(line, require_dollars=True, last=True)
        desc = re.sub(r"\s*-?\$\s*[0-9,]+\.[0-9]{2}.*$", "", line).strip(" :-|•")
        if amount is not None and len(desc) >= 3:
            qty = None
            qm = re.search(r"(?:x|×)\s*(\d+(?:\.\d+)?)\b", desc, re.I)
            if qm:
                qty = float(qm.group(1))
                desc = re.sub(r"\s*(?:x|×)\s*\d+(?:\.\d+)?\b", "", desc, flags=re.I).strip()
            norm_desc = re.sub(r"[^a-z0-9]+", " ", desc.lower()).strip()
            key = (norm_desc, qty, amount)
            if key in seen_items:
                continue
            # OCR variants can repeat the same item from different preprocessing runs.
            duplicate_idx = None
            for idx, existing in enumerate(item_lines):
                if existing.get("quantity") == qty and existing.get("line_total") == amount:
                    ex_norm = re.sub(r"[^a-z0-9]+", " ", existing["description"].lower()).strip()
                    if difflib.SequenceMatcher(None, norm_desc, ex_norm).ratio() >= 0.92:
                        duplicate_idx = idx
                        break
            if duplicate_idx is not None:
                if len(desc) > len(item_lines[duplicate_idx]["description"]):
                    item_lines[duplicate_idx]["description"] = desc[:200]
                continue
            seen_items.add(key)
            item_lines.append({"description": desc[:200], "quantity": qty, "line_total": amount})
    fields["items"] = item_lines[:50]
    return fields


def connect() -> sqlite3.Connection:
    con = sqlite3.connect(DB)
    con.row_factory = sqlite3.Row
    return con


def existing_by_sha(con: sqlite3.Connection, sha: str) -> sqlite3.Row | None:
    return con.execute("SELECT * FROM receipts WHERE image_sha256=?", (sha,)).fetchone()


def insert_or_update(con: sqlite3.Connection, *, source: Path, stored: Path, sha: str, parsed: dict[str, Any], text: str, ocr_info: dict[str, Any], args: argparse.Namespace) -> int:
    captured = args.captured_at or now_local().isoformat(timespec="seconds")
    confidence = "Local OCR via Tesseract. Parsed fields are heuristic; verify important warranty/tax data against original image."
    ocr_json = json.dumps({"parsed": parsed, "ocr": ocr_info}, ensure_ascii=False)

    values = {
        "category": args.category,
        "doc_type": args.doc_type,
        "merchant_name": parsed.get("merchant_name"),
        "receipt_date": parsed.get("receipt_date"),
        "docket_number": parsed.get("docket_number"),
        "invoice_number": parsed.get("invoice_number"),
        "transaction_reference": parsed.get("transaction_reference"),
        "payment_method": parsed.get("payment_method"),
        "subtotal": parsed.get("subtotal"),
        "gst": parsed.get("gst"),
        "total": parsed.get("total"),
        "currency": parsed.get("currency") or "AUD",
        "image_path": str(stored),
        "source_media_path": str(source),
        "image_sha256": sha,
        "ocr_raw_text": text,
        "ocr_json": ocr_json,
        "confidence_notes": confidence,
        "captured_at": captured,
        "notes": args.notes,
    }

    if args.receipt_id:
        rid = args.receipt_id
        assignments = ", ".join(f"{k}=?" for k in values.keys()) + ", updated_at=CURRENT_TIMESTAMP"
        con.execute(f"UPDATE receipts SET {assignments} WHERE id=?", [*values.values(), rid])
    else:
        row = existing_by_sha(con, sha)
        if row:
            rid = int(row["id"])
            if not args.update_existing:
                return rid
            assignments = ", ".join(f"{k}=?" for k in values.keys()) + ", updated_at=CURRENT_TIMESTAMP"
            con.execute(f"UPDATE receipts SET {assignments} WHERE id=?", [*values.values(), rid])
        else:
            cols = ", ".join(values.keys())
            qs = ", ".join("?" for _ in values)
            cur = con.execute(f"INSERT INTO receipts ({cols}) VALUES ({qs})", list(values.values()))
            rid = int(cur.lastrowid)

    if args.replace_items:
        con.execute("DELETE FROM receipt_items WHERE receipt_id=?", (rid,))
        for item in parsed.get("items", []):
            con.execute(
                "INSERT INTO receipt_items (receipt_id, description, quantity, unit_price, line_total, notes) VALUES (?, ?, ?, ?, ?, ?)",
                (rid, item.get("description"), item.get("quantity"), item.get("unit_price"), item.get("line_total"), "auto-parsed from OCR"),
            )
    con.commit()
    return rid


def main() -> None:
    ap = argparse.ArgumentParser(description="Import/OCR a receipt image/PDF into knowledge/receipts")
    ap.add_argument("source", help="Path to receipt image/PDF")
    ap.add_argument("--category", default="Uncategorised")
    ap.add_argument("--doc-type", default="receipt")
    ap.add_argument("--captured-at", help="ISO timestamp; default now")
    ap.add_argument("--receipt-id", type=int, help="Update an existing receipts.id instead of inserting")
    ap.add_argument("--notes", help="Receipt-level note from Michael/Soren, separate from OCR confidence notes")
    ap.add_argument("--update-existing", action="store_true", help="If SHA already exists, update that row")
    ap.add_argument("--replace-items", action=argparse.BooleanOptionalAction, default=True)
    ap.add_argument("--dry-run", action="store_true", help="OCR/parse only; no DB write/copy")
    ap.add_argument("--json", action="store_true", help="Print machine-readable result")
    args = ap.parse_args()

    src = Path(args.source).expanduser().resolve()
    if not src.exists():
        raise SystemExit(f"source not found: {src}")
    if not DB.exists():
        raise SystemExit(f"database not found: {DB}")

    ext = src.suffix.lower()
    sha = sha256_file(src)
    captured = dt.datetime.fromisoformat(args.captured_at) if args.captured_at else now_local()
    ocr_dir = OCR_ROOT / sha[:16]

    if ext in IMAGE_EXTS:
        ocr_info = ocr_image(src, ocr_dir)
    elif ext in PDF_EXTS:
        ocr_info = ocr_pdf(src, ocr_dir)
    elif ext in TEXT_EXTS:
        ocr_dir.mkdir(parents=True, exist_ok=True)
        ocr_info = {"text": clean_text(src.read_text(errors="replace")), "runs": [], "dir": str(ocr_dir)}
        (ocr_dir / "ocr_raw.txt").write_text(ocr_info["text"])
    else:
        raise SystemExit(f"unsupported file type: {ext}")

    text = ocr_info["text"]
    parsed = parse_fields(text)
    result: dict[str, Any] = {"source": str(src), "sha256": sha, "ocr_dir": str(ocr_dir), "parsed": parsed}

    if args.dry_run:
        result["dry_run"] = True
    else:
        try:
            # If the source is already inside the receipts image archive, don't duplicate it.
            src.relative_to(IMAGES.resolve())
            stored = src.relative_to(Path.cwd()) if src.is_relative_to(Path.cwd()) else src
        except Exception:
            stored = copy_original(src, category=args.category, captured=captured, sha=sha)
        with connect() as con:
            rid = insert_or_update(con, source=src, stored=stored, sha=sha, parsed=parsed, text=text, ocr_info=ocr_info, args=args)
        result.update({"receipt_id": rid, "image_path": str(stored)})

    if args.json:
        print(json.dumps(result, indent=2, ensure_ascii=False))
    else:
        print(f"OCR dir: {result['ocr_dir']}")
        if "receipt_id" in result:
            print(f"Receipt ID: {result['receipt_id']}")
            print(f"Image: {result['image_path']}")
        print("Parsed:")
        for k, v in parsed.items():
            if k == "items":
                print(f"  items: {len(v)}")
                for item in v[:10]:
                    print(f"    - {item}")
            else:
                print(f"  {k}: {v}")


if __name__ == "__main__":
    main()