#!/usr/bin/env python3
"""Batch-convert a tree of text PDFs to markdown for whole-file LLM reading.

Walks a source directory (recursing into topic subfolders), converts each PDF to
markdown with pymupdf4llm, and writes a mirrored .md tree under the output dir.

Design notes:
- Idempotent: a PDF is skipped when its .md already exists and is newer.
- No-text-layer PDFs (scans) yield ~no extractable text. Those are detected
  cheaply *before* the expensive markdown pass, logged to needs-ocr.txt, and
  left unwritten (no empty markdown) so they can go through an OCR tool later.
- Output is optimized for a model reading the entire file: a small provenance
  header, then page-separated markdown that keeps headings / lists / tables.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import pymupdf  # bundled with pymupdf4llm
import pymupdf4llm

# A scan (image-only PDF) extracts almost nothing. Real text PDFs comfortably
# exceed this; the threshold only has to separate "basically empty" from "has a
# text layer", so it is deliberately low to avoid false OCR flags.
MIN_CHARS_PER_PAGE = 50


def extractable_chars(pdf_path: Path) -> tuple[int, int]:
    """Return (total_text_chars, page_count) using a cheap text extraction."""
    with pymupdf.open(pdf_path) as doc:
        pages = doc.page_count
        total = sum(len(page.get_text("text")) for page in doc)
    return total, pages


def has_text_layer(total_chars: int, pages: int) -> bool:
    if pages == 0:
        return False
    return (total_chars / pages) >= MIN_CHARS_PER_PAGE


def md_is_current(pdf_path: Path, md_path: Path) -> bool:
    """True when md exists and is at least as new as the source PDF."""
    return md_path.exists() and md_path.stat().st_mtime >= pdf_path.stat().st_mtime


def plain_text_markdown(pdf_path: Path) -> str:
    """Fallback: raw per-page text with the same '-----' page separators.

    Loses heading/list/table structure but never loses the text itself."""
    parts = []
    with pymupdf.open(pdf_path) as doc:
        for page in doc:
            parts.append(page.get_text("text").strip())
    return "\n\n-----\n\n".join(parts)


def convert_one(pdf_path: Path, md_path: Path, raw_chars: int) -> str:
    """Write markdown for one PDF. Returns the method used: 'pymupdf4llm' or 'plain'.

    pymupdf4llm gives structured markdown ("-----" page separators preserve page
    boundaries for citation). On some PDFs its layout pass emits almost nothing
    even though a text layer exists; when its output is implausibly small versus
    the raw extractable text, fall back to plain text extraction.
    """
    md_path.parent.mkdir(parents=True, exist_ok=True)
    body = pymupdf4llm.to_markdown(str(pdf_path), write_images=False, show_progress=False)
    method = "pymupdf4llm"
    # Generous floor: only trips when pymupdf4llm essentially gave up.
    if len(body.strip()) < 0.2 * raw_chars:
        body = plain_text_markdown(pdf_path)
        method = "plain"
    header = f"# {pdf_path.stem}\n\n> Source: `{pdf_path.name}`\n\n---\n\n"
    md_path.write_text(header + body, encoding="utf-8")
    return method


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--src", type=Path, default=Path("pdfs"), help="source dir of PDFs (default: pdfs)")
    parser.add_argument("--out", type=Path, default=Path("md"), help="output dir for markdown (default: md)")
    parser.add_argument(
        "--ocr-list", type=Path, default=Path("needs-ocr.txt"),
        help="file listing PDFs with no usable text layer (default: needs-ocr.txt)",
    )
    parser.add_argument("--force", action="store_true", help="reconvert even if md is current")
    args = parser.parse_args(argv)

    if not args.src.is_dir():
        print(f"error: source dir not found: {args.src}", file=sys.stderr)
        return 2

    pdfs = sorted(args.src.rglob("*.pdf"))
    if not pdfs:
        print(f"No PDFs found under {args.src}")
        return 0

    converted, skipped, failed, plain = 0, 0, 0, 0
    needs_ocr: list[Path] = []

    for pdf in pdfs:
        rel = pdf.relative_to(args.src).with_suffix(".md")
        md_path = args.out / rel

        if not args.force and md_is_current(pdf, md_path):
            print(f"skip   {rel}  (up to date)")
            skipped += 1
            continue

        try:
            total, pages = extractable_chars(pdf)
        except Exception as exc:  # unreadable / corrupt PDF
            print(f"FAIL   {rel}  ({type(exc).__name__}: {exc})")
            failed += 1
            continue

        if not has_text_layer(total, pages):
            print(f"ocr?   {rel}  ({total} chars / {pages} pages -> no text layer)")
            needs_ocr.append(pdf)
            continue

        try:
            method = convert_one(pdf, md_path, total)
        except Exception as exc:
            print(f"FAIL   {rel}  ({type(exc).__name__}: {exc})")
            failed += 1
            continue

        tag = "  [plain-text fallback]" if method == "plain" else ""
        if method == "plain":
            plain += 1
        print(f"conv   {rel}  ({pages} pages){tag}")
        converted += 1

    # Rewritten fresh each run: flagged PDFs are re-checked every time since they
    # never produce markdown to skip on.
    if needs_ocr:
        args.ocr_list.write_text("\n".join(str(p) for p in needs_ocr) + "\n", encoding="utf-8")
    elif args.ocr_list.exists():
        args.ocr_list.unlink()

    print("\n" + "=" * 48)
    print(f"converted:       {converted}" + (f"  ({plain} via plain-text fallback)" if plain else ""))
    print(f"skipped:         {skipped}")
    print(f"flagged for OCR: {len(needs_ocr)}")
    if failed:
        print(f"failed:          {failed}")
    if needs_ocr:
        print(f"\nSee {args.ocr_list} -> run these through marker/docling (OCR).")
    return 1 if failed else 0


if __name__ == "__main__":
    raise SystemExit(main())