"""
Split multi-invoice PDFs into one PDF per invoice.

For each PDF in invoices/ (skipping _non_jei/):
  1. Detect distinct invoice numbers across pages (regex on extracted text,
     with OCR fallback from data/review/ocr/<stem>/page_NNN.txt)
  2. If 0 or 1 distinct invoice number found, leave the file alone —
     rename_invoices.py handles single-invoice naming.
  3. If 2+ distinct invoice numbers found:
     a. Group consecutive pages by invoice number (so a multi-page proforma
        like 151779 with 6 pages of the same invoice stays as ONE output PDF)
     b. Write each group to invoices/<invoice_no>.pdf via PyMuPDF
     c. Move the source PDF to invoices/_non_jei/

Idempotent: skips source PDFs that already have no invoice number
detected, or where every output target already exists and has the same
page count (avoids re-doing work).

Usage:
    python scripts/split_invoices.py           # process all multi-invoice PDFs
    python scripts/split_invoices.py <pdf>...  # process specific files only

Notes on edge cases:
- A "multi-page proforma" (e.g. PROFORMA 151779 over 6 pages) appears as a
  single invoice_no across all pages. The script keeps those as ONE output.
- A "multi-invoice bundle" (e.g. 1.Invoice_29.05.pdf with 4 invoices on
  separate pages) is split into 4 separate output PDFs.
- Scanned PDFs need ocr_invoices.py to have been run first so the OCR'd text
  exists in data/review/ocr/. Otherwise we can't detect the invoice numbers.
"""
import os, re, shutil, sys
from pathlib import Path
import pdfplumber
import fitz  # PyMuPDF — used to write new PDFs

ROOT = Path(__file__).resolve().parent.parent
INVOICES = ROOT / "invoices"
NON_JEI = INVOICES / "_non_jei"
OCR_DIR = ROOT / "data" / "review" / "ocr"

NON_JEI.mkdir(exist_ok=True)

INVOICE_NO_RE = re.compile(r"(?:Invoice|Quotation)\s*No\s*(\d{3,7})", re.IGNORECASE)


def page_invoice_no(pdf_path: Path, page_index: int) -> str | None:
    """Return invoice_no found on a given page (0-indexed), with OCR fallback."""
    try:
        with pdfplumber.open(pdf_path) as pdf:
            text = pdf.pages[page_index].extract_text() or ""
    except Exception:
        text = ""
    m = INVOICE_NO_RE.search(text)
    if m:
        return m.group(1)
    # OCR fallback
    ocr_file = OCR_DIR / pdf_path.stem.strip() / f"page_{page_index + 1:03d}.txt"
    if ocr_file.exists():
        text = ocr_file.read_text(encoding="utf-8", errors="replace")
        m = INVOICE_NO_RE.search(text)
        if m:
            return m.group(1)
    return None


def page_count(pdf_path: Path) -> int:
    try:
        with pdfplumber.open(pdf_path) as pdf:
            return len(pdf.pages)
    except Exception:
        return 0


def split_pdf(src: Path, dry_run: bool = False) -> dict:
    """Split a multi-invoice PDF into per-invoice PDFs. Returns a summary dict."""
    n_pages = page_count(src)
    if n_pages == 0:
        return {"src": src.name, "status": "unreadable"}

    # Step 1: per-page invoice number
    per_page = []
    for i in range(n_pages):
        per_page.append(page_invoice_no(src, i))

    # Drop trailing/leading None pages (cover sheets, blanks)
    distinct = set(per_page) - {None}
    if len(distinct) <= 1:
        return {"src": src.name, "status": "single-invoice", "pages": n_pages,
                "invoice_no": next(iter(distinct)) if distinct else None}

    # Step 2: group consecutive pages by invoice_no
    groups = []
    cur_no = None
    cur_pages = []
    for idx, no in enumerate(per_page):
        if no == cur_no:
            cur_pages.append(idx)
        else:
            if cur_pages and cur_no:
                groups.append((cur_no, cur_pages))
            cur_no = no
            cur_pages = [idx] if no else []
    if cur_pages and cur_no:
        groups.append((cur_no, cur_pages))

    # Step 3: write each group
    written = []
    skipped_existing = []
    fitz_doc = fitz.open(str(src))
    try:
        for inv_no, page_idxs in groups:
            target = INVOICES / f"{inv_no}.pdf"
            if target.exists():
                existing_pages = page_count(target)
                if existing_pages == len(page_idxs):
                    skipped_existing.append(inv_no)
                    continue
                else:
                    # collision with a DIFFERENT existing file of the same name
                    # — append a suffix so we never overwrite
                    i = 2
                    while True:
                        alt = INVOICES / f"{inv_no}_{i}.pdf"
                        if not alt.exists():
                            target = alt
                            break
                        i += 1
            if dry_run:
                written.append(target.name)
                continue
            new_doc = fitz.open()
            for p in page_idxs:
                new_doc.insert_pdf(fitz_doc, from_page=p, to_page=p)
            new_doc.save(str(target))
            new_doc.close()
            written.append(target.name)
    finally:
        fitz_doc.close()

    # Step 4: move source to _non_jei/ (only if all groups were written)
    if not dry_run and written:
        dest = NON_JEI / src.name
        if dest.exists():
            dest.unlink()
        shutil.move(str(src), str(dest))

    return {"src": src.name, "status": "split", "groups": groups,
            "written": written, "skipped_existing": skipped_existing,
            "moved_to": "_non_jei/" if (not dry_run and written) else None}


def main():
    args = sys.argv[1:]
    if args:
        # Process specified files only
        pdfs = [Path(a) for a in args]
    else:
        pdfs = sorted(p for p in INVOICES.glob("*.pdf") if "_non_jei" not in p.parts)

    if not pdfs:
        print("No PDFs to process")
        return

    print(f"Scanning {len(pdfs)} PDF(s) in invoices/\n")
    summary = []
    for p in pdfs:
        if not p.exists():
            print(f"  MISSING: {p}")
            continue
        result = split_pdf(p)
        summary.append(result)
        if result["status"] == "single-invoice":
            print(f"  KEEP   {p.name:40s}  (single-invoice{', invoice_no=' + result['invoice_no'] if result.get('invoice_no') else ', unparseable'})")
        elif result["status"] == "unreadable":
            print(f"  SKIP   {p.name:40s}  (unreadable)")
        elif result["status"] == "split":
            inv_nos = [g[0] for g in result["groups"]]
            print(f"  SPLIT  {p.name:40s}  -> {len(result['groups'])} files: {inv_nos}")
            for w in result["written"]:
                print(f"           + {w}")
            for s in result["skipped_existing"]:
                print(f"           = {s}.pdf (already exists, same page count)")

    # Final state
    print(f"\nFinal state of invoices/:")
    kept = sorted(p for p in INVOICES.glob("*.pdf"))
    for p in kept:
        print(f"  {p.name}")
    print(f"\nTotal in invoices/: {len(kept)}")


if __name__ == "__main__":
    main()