"""
OCR scanned PDFs using Tesseract + PyMuPDF (no poppler needed).
Reads each unparseable PDF page-by-page, renders to image, OCRs the image,
and saves the text to data/review/ocr/<invoice_no>.txt.

Then merges with the text-based extraction by writing OCR'd text back into
a fake pdfplumber-compatible page object via monkey-patching.

Run: python scripts/ocr_invoices.py
"""
import os, sys, subprocess, json
from pathlib import Path
import fitz  # PyMuPDF

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

# Find tesseract
TESSERACT = None
for p in ["C:/Program Files/Tesseract-OCR/tesseract.exe", "C:/Program Files (x86)/Tesseract-OCR/tesseract.exe"]:
    if os.path.exists(p):
        TESSERACT = p
        break
if not TESSERACT:
    sys.exit("tesseract.exe not found - install via winget install UB-Mannheim.TesseractOCR")

OCR_DIR.mkdir(parents=True, exist_ok=True)


def ocr_page_image(img_bytes: bytes) -> str:
    """OCR a single PNG image, return extracted text."""
    # Write to temp file
    tmp_in = OCR_DIR / "_tmp_in.png"
    tmp_out = OCR_DIR / "_tmp_out"
    tmp_in.write_bytes(img_bytes)
    try:
        result = subprocess.run(
            [TESSERACT, str(tmp_in), str(tmp_out), "-l", "eng"],
            capture_output=True, text=True, timeout=60,
        )
        out_file = tmp_out.with_suffix(".txt")
        if out_file.exists():
            return out_file.read_text(encoding="utf-8", errors="replace")
        return ""
    finally:
        tmp_in.unlink(missing_ok=True)
        (tmp_out.with_suffix(".txt")).unlink(missing_ok=True)


def ocr_pdf(pdf_path: Path) -> list[str]:
    """Return list of OCR'd text per page."""
    texts = []
    with fitz.open(pdf_path) as doc:
        for page_num, page in enumerate(doc, start=1):
            # Render at 3x for better OCR accuracy
            mat = fitz.Matrix(3, 3)
            pix = page.get_pixmap(matrix=mat)
            img_bytes = pix.tobytes("png")
            text = ocr_page_image(img_bytes)
            texts.append(text)
            print(f"  {pdf_path.name} page {page_num}/{len(doc)}: {len(text)} chars")
    return texts


def main():
    pdfs = sorted(p for p in INVOICES.rglob("*.pdf") if "_non_jei" not in p.parts)
    errors = 0
    for pdf in pdfs:
        try:
            out_dir = OCR_DIR / pdf.stem.strip()
            out_dir.mkdir(exist_ok=True)
            # OCR all pages
            texts = ocr_pdf(pdf)
            for i, text in enumerate(texts, start=1):
                (out_dir / f"page_{i:03d}.txt").write_text(text, encoding="utf-8")
            # Combined file
            combined = "\n=== PAGE BREAK ===\n".join(texts)
            (OCR_DIR / f"{pdf.stem.strip()}.txt").write_text(combined, encoding="utf-8")
        except Exception as e:
            errors += 1
            print(f"  ERR {pdf.name}: {e}")
    print(f"\nDone. OCR text in {OCR_DIR}. Errors: {errors}")


if __name__ == "__main__":
    main()