pdf_chunker.py

import json
import os
import sys
from pathlib import Path

try:
    from pypdf import PdfReader, PdfWriter
except ImportError:
    print(json.dumps({"error": "pypdf no instalado. Ejecuta: pip install pypdf"}))
    sys.exit(1)


def chunk_pdf(pdf_path: str, pages_per_chunk: int = 40, threshold: int = 50) -> dict:
    src = Path(pdf_path)
    if not src.exists():
        return {"error": f"No existe el PDF: {pdf_path}"}

    reader = PdfReader(str(src))
    total_pages = len(reader.pages)

    # PDF pequeno: no se divide, se devuelve el original como unico "chunk".
    if total_pages <= threshold:
        return {
            "total_pages": total_pages,
            "total_chunks": 1,
            "chunked": False,
            "chunks": [{
                "idx": 0,
                "path": str(src.resolve()),
                "from": 1,
                "to": total_pages,
            }],
        }

    # PDF grande: dividir.
    stem = src.stem  # p.ej. "estado_abc123"
    out_dir = src.parent
    chunks = []
    chunk_idx = 0
    for start in range(0, total_pages, pages_per_chunk):
        end = min(start + pages_per_chunk, total_pages)
        writer = PdfWriter()
        for p in range(start, end):
            writer.add_page(reader.pages[p])
        chunk_path = out_dir / f"{stem}_chunk_{chunk_idx:02d}.pdf"
        with open(chunk_path, "wb") as f:
            writer.write(f)
        chunks.append({
            "idx": chunk_idx,
            "path": str(chunk_path.resolve()),
            "from": start + 1,
            "to": end,
        })
        chunk_idx += 1

    return {
        "total_pages": total_pages,
        "total_chunks": len(chunks),
        "chunked": True,
        "chunks": chunks,
    }


def main():
    if len(sys.argv) < 2:
        print(json.dumps({"error": "Uso: pdf_chunker.py <pdf_path> [pages_per_chunk] [threshold]"}))
        sys.exit(1)

    pdf_path = sys.argv[1]
    pages_per_chunk = int(sys.argv[2]) if len(sys.argv) > 2 else 40
    threshold = int(sys.argv[3]) if len(sys.argv) > 3 else 50

    result = chunk_pdf(pdf_path, pages_per_chunk, threshold)
    print(json.dumps(result, ensure_ascii=False))


if __name__ == "__main__":
    main()