src/knowledge_parsing/parse.py

"""Tahap 2 — parsing dokumen dengan MinerU, dengan cache berbasis isi.

Kenapa cache-nya dikunci ke ISI dokumen, bukan ke nomor run: parsing adalah
tahap termahal (19 dtk/halaman di CPU). Kalau cache-nya per run, tiap kali
normalizer diubah dan pipeline dijalankan ulang, parsing ikut terulang percuma.
Dengan kunci isi+pengaturan+versi, dokumen yang sama hanya diparse sekali —
mau normalizer diubah 20 kali sekalipun.

Kunci cache = sha256(isi berkas) + sidik jari pengaturan + versi MinerU.
Versi ikut supaya angka hasil MinerU lama tidak diam-diam tercampur dengan
hasil versi baru — ini penting karena angkanya dipakai untuk justifikasi biaya.
"""

from __future__ import annotations

import hashlib
import json
import shutil
import time
from dataclasses import dataclass
from pathlib import Path

from .config import PipelineConfig


@dataclass
class ParseResult:
    doc_id: str
    source: Path
    cache_dir: Path          # folder berisi output MinerU apa adanya
    content_list: Path       # *_content_list.json
    middle_json: Path | None
    pages: int
    seconds: float
    from_cache: bool
    backend_tercatat: str | None    # dibaca dari _middle.json, BUKAN dari config
    mineru_version: str | None


def versi_mineru() -> str:
    try:
        from mineru.version import __version__  # type: ignore
        return str(__version__)
    except Exception:
        try:
            from importlib.metadata import version
            return version("mineru")
        except Exception:
            return "unknown"


def hash_berkas(path: Path, potong: int = 16) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for blok in iter(lambda: f.read(1 << 20), b""):
            h.update(blok)
    return h.hexdigest()[:potong]


def kunci_cache(source: Path, cfg: PipelineConfig) -> str:
    return f"{hash_berkas(source)}-{cfg.fingerprint()}-{versi_mineru()}"


def _cari_output(root: Path) -> tuple[Path | None, Path | None]:
    """Cari content_list & middle json di bawah root.

    Sengaja pakai glob, bukan path tetap: struktur subfolder MinerU berbeda
    antar backend ('auto' untuk pipeline, lain untuk vlm). Glob membuat modul
    ini tidak perlu diubah saat backend diganti.
    """
    content = next(iter(sorted(root.rglob("*_content_list.json"))), None)
    middle = next(iter(sorted(root.rglob("*_middle.json"))), None)
    return content, middle


def _baca_middle(middle: Path | None) -> tuple[str | None, str | None]:
    """Ambil backend & versi yang BENAR-BENAR dipakai, dari output MinerU."""
    if not middle or not middle.exists():
        return None, None
    try:
        d = json.loads(middle.read_text(encoding="utf-8"))
        return d.get("_backend"), d.get("_version_name")
    except Exception:
        return None, None


def parse_dokumen(source: Path, cfg: PipelineConfig) -> ParseResult:
    """Parse satu dokumen. Kalau sudah ada di cache, tidak dijalankan ulang."""
    doc_id = source.stem
    tujuan = cfg.cache_dir / "parse" / kunci_cache(source, cfg)
    penanda = tujuan / ".selesai"

    if penanda.exists():
        content, middle = _cari_output(tujuan)
        if content:
            meta = json.loads(penanda.read_text(encoding="utf-8"))
            backend_tercatat, versi = _baca_middle(middle)
            return ParseResult(
                doc_id=doc_id, source=source, cache_dir=tujuan,
                content_list=content, middle_json=middle,
                pages=meta.get("pages", 0), seconds=meta.get("seconds", 0.0),
                from_cache=True,
                backend_tercatat=backend_tercatat, mineru_version=versi,
            )
        shutil.rmtree(tujuan, ignore_errors=True)   # cache rusak, ulangi

    from mineru.cli.common import do_parse, read_fn

    # Playground melakukan ini sebelum parsing; alias seperti "latin"/"en"
    # diterjemahkan ke kode kanonik. Disamakan supaya hasilnya identik.
    lang = cfg.lang_kanonik()

    sedang = tujuan.with_suffix(".sedang")
    shutil.rmtree(sedang, ignore_errors=True)
    sedang.mkdir(parents=True, exist_ok=True)

    pdf_bytes = read_fn(source)
    mulai = time.perf_counter()
    do_parse(
        output_dir=str(sedang),
        pdf_file_names=[doc_id],
        pdf_bytes_list=[pdf_bytes],
        p_lang_list=[lang],
        backend=cfg.backend,
        formula_enable=cfg.formula_enable,
        table_enable=cfg.table_enable,
        f_draw_layout_bbox=cfg.write_debug_pdf,
        f_draw_span_bbox=cfg.write_debug_pdf,
        start_page_id=cfg.start_page,
        end_page_id=cfg.end_page,
        **({"effort": cfg.effort} if cfg.backend == "hybrid" else {}),
    )
    detik = time.perf_counter() - mulai

    content, middle = _cari_output(sedang)
    if content is None:
        raise RuntimeError(
            f"MinerU selesai tapi *_content_list.json tidak ditemukan di {sedang}"
        )

    halaman = _hitung_halaman(source, cfg)
    sedang.rename(tujuan)
    content = tujuan / content.relative_to(sedang)
    middle = tujuan / middle.relative_to(sedang) if middle else None

    penanda.write_text(
        json.dumps({"pages": halaman, "seconds": detik, "doc_id": doc_id}),
        encoding="utf-8",
    )
    backend_tercatat, versi = _baca_middle(middle)

    return ParseResult(
        doc_id=doc_id, source=source, cache_dir=tujuan,
        content_list=content, middle_json=middle,
        pages=halaman, seconds=detik, from_cache=False,
        backend_tercatat=backend_tercatat, mineru_version=versi,
    )


def _hitung_halaman(source: Path, cfg: PipelineConfig) -> int:
    if cfg.end_page is not None:
        return cfg.end_page - cfg.start_page + 1
    try:
        from pypdf import PdfReader
        return len(PdfReader(str(source)).pages)
    except Exception:
        return 0