src/autocad_bench/bundle.py

"""Portable, integrity-bound benchmark data bundles."""

from __future__ import annotations

import hashlib
import os
import re
import shutil
import tempfile
from pathlib import Path, PurePosixPath
from typing import Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator

from autocad_bench.common.io import atomic_write_json
from autocad_bench.common.paths import BENCHMARK_ROOT
from autocad_bench.evaluation.scoring.gold_cache import GoldCacheStore
from autocad_bench.tasks.audit import load_gold_audit
from autocad_bench.tasks.manifest import load_manifest
from autocad_bench.tasks.suite import load_benchmark_task_ids

BUNDLE_FILENAME = "autocad-bench-bundle.json"
BUNDLE_SCHEMA_VERSION = 1
BUNDLE_VERSION = "autocad-bench-corpus-v1"


class BundleError(RuntimeError):
    """A benchmark bundle is incomplete, unsafe, or not reproducible."""


def _sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


class BundleFile(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    path: str
    bytes: int = Field(ge=0)
    sha256: str = Field(pattern=r"^[a-f0-9]{64}$")

    @field_validator("path")
    @classmethod
    def validate_path(cls, value: str) -> str:
        path = PurePosixPath(value)
        if (
            path.is_absolute()
            or not path.parts
            or ".." in path.parts
            or value != path.as_posix()
        ):
            raise ValueError("bundle file paths must be normalized and relative")
        if value == BUNDLE_FILENAME:
            raise ValueError("bundle manifest cannot list itself")
        return value


class BenchmarkBundle(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    schema_version: Literal[1] = BUNDLE_SCHEMA_VERSION
    bundle_version: Literal["autocad-bench-corpus-v1"] = BUNDLE_VERSION
    task_count: int = Field(gt=0)
    evaluator_versions: tuple[str, ...] = ()
    files: tuple[BundleFile, ...] = Field(min_length=1)

    @field_validator("files")
    @classmethod
    def validate_unique_files(
        cls,
        value: tuple[BundleFile, ...],
    ) -> tuple[BundleFile, ...]:
        paths = [entry.path for entry in value]
        if len(paths) != len(set(paths)):
            raise ValueError("bundle contains duplicate file paths")
        if paths != sorted(paths):
            raise ValueError("bundle files must be sorted by path")
        return value

    @field_validator("evaluator_versions")
    @classmethod
    def validate_unique_versions(cls, value: tuple[str, ...]) -> tuple[str, ...]:
        if tuple(sorted(set(value))) != value:
            raise ValueError("evaluator_versions must be sorted and unique")
        if any(not re.fullmatch(r"[A-Za-z0-9_.-]+", item) for item in value):
            raise ValueError("evaluator_versions contains an unsafe value")
        return value


def _source_files(
    root: Path,
    *,
    evaluator_versions: tuple[str, ...],
) -> tuple[Path, ...]:
    entries = load_manifest(root=root)
    load_gold_audit(root=root)
    load_benchmark_task_ids(root=root)

    relative_paths = {
        Path("tasks/README.md"),
        Path("tasks/manifest.jsonl"),
        Path("tasks/gold-audit.jsonl"),
        Path("tasks/benchmark-suite-50.txt"),
    }
    for entry in entries:
        relative_paths.add(Path(entry.image_path))
        relative_paths.add(Path(entry.gold_path))

    for evaluator_version in evaluator_versions:
        store = GoldCacheStore(
            evaluator_version=evaluator_version,
            root=root / "artifacts" / "gold-cache",
        )
        expected_checksums = {
            entry.task_id: _sha256_file(entry.resolve_gold_path(root=root))
            for entry in entries
        }
        try:
            store.verify(
                [entry.task_id for entry in entries],
                expected_source_sha256=expected_checksums,
            )
        except Exception as exc:
            raise BundleError(
                f"gold cache is not release-ready for {evaluator_version}: {exc}"
            ) from exc
        for entry in entries:
            base = (
                Path("artifacts")
                / "gold-cache"
                / evaluator_version
                / entry.task_id
            )
            relative_paths.add(base / "metadata.json")
            relative_paths.add(base / "render.png")

    paths: list[Path] = []
    for relative in sorted(relative_paths, key=lambda item: item.as_posix()):
        source = (root / relative).resolve()
        if not source.is_relative_to(root):
            raise BundleError(f"bundle source escapes root: {relative.as_posix()}")
        if source.is_symlink() or not source.is_file():
            raise BundleError(
                f"bundle source must be a regular file: {relative.as_posix()}"
            )
        paths.append(source)
    return tuple(paths)


def export_bundle(
    output_root: Path | str,
    *,
    source_root: Path | str = BENCHMARK_ROOT,
    evaluator_versions: tuple[str, ...] = (),
) -> BenchmarkBundle:
    """Export an allowlisted benchmark directory and bind every byte."""

    source = Path(source_root).expanduser().resolve()
    output = Path(output_root).expanduser().resolve()
    if output.exists():
        raise BundleError(f"bundle output already exists: {output}")
    output.parent.mkdir(parents=True, exist_ok=True)

    versions = tuple(sorted(set(evaluator_versions)))
    files = _source_files(source, evaluator_versions=versions)
    staging = Path(
        tempfile.mkdtemp(
            prefix=f".{output.name}.staging-",
            dir=output.parent,
        )
    )
    try:
        manifest_files: list[BundleFile] = []
        for source_path in files:
            relative = source_path.relative_to(source)
            destination = staging / relative
            destination.parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(source_path, destination)
            manifest_files.append(
                BundleFile(
                    path=relative.as_posix(),
                    bytes=destination.stat().st_size,
                    sha256=_sha256_file(destination),
                )
            )
        bundle = BenchmarkBundle(
            task_count=len(load_manifest(root=source)),
            evaluator_versions=versions,
            files=tuple(manifest_files),
        )
        atomic_write_json(
            staging / BUNDLE_FILENAME,
            bundle.model_dump(mode="json"),
        )
        os.replace(staging, output)
        return bundle
    except BaseException:
        shutil.rmtree(staging, ignore_errors=True)
        raise


def validate_bundle(root: Path | str) -> BenchmarkBundle:
    """Verify byte integrity and the semantic task/evaluator contracts."""

    bundle_root = Path(root).expanduser().resolve()
    manifest_path = bundle_root / BUNDLE_FILENAME
    if not manifest_path.is_file() or manifest_path.is_symlink():
        raise BundleError(f"bundle manifest not found: {manifest_path}")
    try:
        bundle = BenchmarkBundle.model_validate_json(
            manifest_path.read_text(encoding="utf-8")
        )
    except Exception as exc:
        raise BundleError(f"invalid bundle manifest: {exc}") from exc

    for candidate in bundle_root.rglob("*"):
        if candidate.is_symlink():
            raise BundleError(
                "bundle must not contain symbolic links: "
                + candidate.relative_to(bundle_root).as_posix()
            )
    expected_paths = {entry.path for entry in bundle.files} | {BUNDLE_FILENAME}
    actual_paths = {
        path.relative_to(bundle_root).as_posix()
        for path in bundle_root.rglob("*")
        if path.is_file()
    }
    if actual_paths != expected_paths:
        missing = sorted(expected_paths - actual_paths)
        unexpected = sorted(actual_paths - expected_paths)
        raise BundleError(
            f"bundle file inventory mismatch: missing={missing}, "
            f"unexpected={unexpected}"
        )

    for entry in bundle.files:
        path = (bundle_root / entry.path).resolve()
        if not path.is_relative_to(bundle_root) or not path.is_file():
            raise BundleError(f"bundle file is unavailable: {entry.path}")
        if path.stat().st_size != entry.bytes:
            raise BundleError(f"bundle file size changed: {entry.path}")
        if _sha256_file(path) != entry.sha256:
            raise BundleError(f"bundle file checksum changed: {entry.path}")

    entries = load_manifest(root=bundle_root)
    if len(entries) != bundle.task_count:
        raise BundleError(
            f"bundle task count changed: expected {bundle.task_count}, "
            f"found {len(entries)}"
        )
    load_gold_audit(root=bundle_root)
    load_benchmark_task_ids(root=bundle_root)
    expected_checksums = {
        entry.task_id: _sha256_file(entry.resolve_gold_path(root=bundle_root))
        for entry in entries
    }
    for evaluator_version in bundle.evaluator_versions:
        store = GoldCacheStore(
            evaluator_version=evaluator_version,
            root=bundle_root / "artifacts" / "gold-cache",
        )
        try:
            store.verify(
                [entry.task_id for entry in entries],
                expected_source_sha256=expected_checksums,
            )
        except Exception as exc:
            raise BundleError(
                f"invalid gold cache for {evaluator_version}: {exc}"
            ) from exc
    return bundle


__all__ = [
    "BUNDLE_FILENAME",
    "BUNDLE_SCHEMA_VERSION",
    "BUNDLE_VERSION",
    "BenchmarkBundle",
    "BundleError",
    "BundleFile",
    "export_bundle",
    "validate_bundle",
]