#!/usr/bin/env python3
"""Populate the gitignored B1 image-only eval corpus from the .94 exam-board store.

The B1 eval papers are NOT committed (third-party copyright; served only via signed URLs).
This script downloads each B1_GEOMETRY paper's `storage_loc` object from cc.examboards via the
Storage API into its local `pdf` path (under samples/b1/), so finalize.py --b1-only and the
B1-2/B1-3 generalization work can run against a real corpus.

Run from api/services/docling/ inside the cc-api-dev container (SUPABASE_URL/SERVICE_ROLE_KEY in env):
    python3 scripts/fetch_b1_corpus.py            # fetch all B1 papers (skip existing)
    python3 scripts/fetch_b1_corpus.py --force    # re-download
    python3 scripts/fetch_b1_corpus.py --only b1-aqa-physics-7408-1-2022jun
    python3 scripts/fetch_b1_corpus.py --list     # show what would be fetched, no download
"""
from __future__ import annotations

import argparse
import os
import sys

# Import the canonical B1 corpus definition (slug, storage_loc, local pdf path) from finalize.
_HERE = os.path.dirname(os.path.abspath(__file__))
_DOCLING_DIR = os.path.dirname(_HERE)
sys.path.insert(0, _DOCLING_DIR)
from finalize import B1_GEOMETRY  # noqa: E402


def _split_storage_loc(storage_loc: str) -> tuple[str, str]:
    """'cc.examboards/aqa/biology/7402/1/2023-jun/qp.pdf' -> ('cc.examboards', 'aqa/.../qp.pdf')."""
    bucket, _, path = storage_loc.partition("/")
    if not bucket or not path:
        raise ValueError(f"malformed storage_loc: {storage_loc!r}")
    return bucket, path


def _entries(only: str | None):
    for p in B1_GEOMETRY:
        loc = p.get("storage_loc")
        pdf = p.get("pdf")
        if not loc or not pdf:
            continue
        if only and p.get("slug") != only:
            continue
        yield p["slug"], loc, pdf


def main() -> int:
    ap = argparse.ArgumentParser(description="Fetch the B1 image-only eval corpus from .94 cc.examboards")
    ap.add_argument("--force", action="store_true", help="re-download even if the local file exists")
    ap.add_argument("--only", help="fetch a single paper by slug")
    ap.add_argument("--list", action="store_true", help="list what would be fetched and exit")
    args = ap.parse_args()

    todo = list(_entries(args.only))
    if not todo:
        print("no matching B1 papers", file=sys.stderr)
        return 1

    if args.list:
        for slug, loc, pdf in todo:
            print(f"{slug}\t{loc}\t-> {pdf}")
        return 0

    from modules.database.supabase.utils.storage import StorageAdmin
    storage = StorageAdmin()

    ok = skipped = 0
    for slug, loc, pdf in todo:
        dest = os.path.join(_DOCLING_DIR, pdf) if not os.path.isabs(pdf) else pdf
        if os.path.exists(dest) and not args.force:
            print(f"[skip] {slug} (exists)")
            skipped += 1
            continue
        bucket, path = _split_storage_loc(loc)
        data = storage.download_file(bucket, path)
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        with open(dest, "wb") as fh:
            fh.write(data)
        print(f"[ok]   {slug} <- {bucket}/{path} ({len(data)} bytes)")
        ok += 1

    print(f"fetched {ok}, skipped {skipped}, of {len(todo)}")
    return 0


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