"""Export original page pixels losslessly, with per-document provenance."""

import json
from pathlib import Path
from time import perf_counter

from visual_doc_rag.data import load_tables, write_json


def export_pages(subset: str, data_dir: Path) -> dict:
    started = perf_counter()
    source, tables = load_tables(subset, data_dir)
    root = data_dir / "pages" / subset / source["revision"]
    root.mkdir(parents=True, exist_ok=True)
    metadata = {row["doc_id"]: row for row in tables["documents_metadata"]}
    pages = []
    for index, row in enumerate(tables["corpus"]):
        path = root / f"{row['corpus_id']}.png"
        with row["image"] as image:
            if not path.exists():
                temporary = path.with_suffix(".tmp")
                image.convert("RGB").save(temporary, format="PNG")
                temporary.replace(path)
            size = list(image.size)
        document = metadata[row["doc_id"]]
        pages.append(
            {
                "corpus_id": row["corpus_id"],
                "doc_id": row["doc_id"],
                "page_number_in_doc": row["page_number_in_doc"],
                "file": path.name,
                "size": size,
                "license": document["license"],
                "source_url": document["url"],
                "source_file_name": document["file_name"],
            }
        )
        if (index + 1) % 100 == 0:
            print(f"Pages: {index + 1}/{len(tables['corpus'])}", flush=True)
    manifest = {
        "dataset": source,
        "pages": pages,
        "format": "PNG, original pixels, RGB",
        "export_seconds": perf_counter() - started,
        "image_bytes": sum((root / row["file"]).stat().st_size for row in pages),
    }
    write_json(root / "manifest.json", manifest)
    return manifest


def load_pages(subset: str, data_dir: Path, revision: str) -> tuple[Path, dict]:
    root = data_dir / "pages" / subset / revision
    path = root / "manifest.json"
    if not path.exists():
        raise ValueError("Page cache missing for this dataset revision. Run `vdr pages` first.")
    return root, json.loads(path.read_text(encoding="utf-8"))