Blame for src/dlb/checkpoints.py

"""Typed checkpoint manifests, safe download helpers, and lock verification."""

from __future__ import annotations

from datetime import datetime, timezone
from fnmatch import fnmatchcase
import hashlib
import json
import os
from pathlib import Path, PurePosixPath
import shutil
import stat
import subprocess
from typing import Literal, Union
from urllib.error import HTTPError
from urllib.request import Request, urlopen

import yaml
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from dlb.io import atomic_json_write, sha256_file


Provenance = Literal["official", "reference_reproduction", "self_trained"]
TeacherFamily = Literal[
    "continuous_flm",
    "continuous_langflow",
    "continuous_rdlm",
    "uniform_duo",
    "masked_mdlm",
    "hybrid_candi",
    "mixed_baselines",
]
TeacherAdapter = Literal["uniform_to_absorbing", "masked_to_absorbing"]


class ManifestModel(BaseModel):
    model_config = ConfigDict(extra="forbid")


class DigestSpec(ManifestModel):
    policy: Literal["sha256", "capture_after_download"]
    sha256: str | None = Field(default=None, pattern=r"^[0-9a-f]{64}$")
    per_file_sha256: dict[str, str] = Field(default_factory=dict)

    @model_validator(mode="after")
    def validate_policy(self) -> "DigestSpec":
        for name, digest in self.per_file_sha256.items():
            safe_remote_path(name)
            if len(digest) != 64 or any(
                character not in "0123456789abcdef" for character in digest
            ):
                raise ValueError(f"invalid per-file SHA-256 for {name}")
        declared = (self.sha256 is not None) + bool(self.per_file_sha256)
        if self.policy == "sha256" and declared != 1:
            raise ValueError("sha256 policy requires one aggregate or per-file digest set")
        if self.policy == "capture_after_download" and declared:
            raise ValueError("capture policy forbids declared SHA-256 values")
        return self


class HuggingFaceSource(ManifestModel):
    repo_id: str = Field(pattern=r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$")
    revision: str = Field(pattern=r"^[0-9a-f]{40}$")
    allow_patterns: list[str] = Field(min_length=1)


class GDriveSource(ManifestModel):
    folder_id: str = Field(pattern=r"^[A-Za-z0-9_-]+$")
    expected_files: dict[str, str] = Field(default_factory=dict)

    @model_validator(mode="after")
    def validate_expected_files(self) -> "GDriveSource":
        if len(self.expected_files.values()) != len(set(self.expected_files.values())):
            raise ValueError("Google Drive file IDs must be unique")
        for path, file_id in self.expected_files.items():
            safe_remote_path(path)
            allowed = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_-"
            if not file_id or any(character not in allowed for character in file_id):
                raise ValueError(f"invalid Google Drive file id for {path}")
        return self


class ZenodoSource(ManifestModel):
    record_id: int = Field(gt=0)
    files: list[str] = Field(min_length=1)
    published_checksums: dict[str, str] = Field(default_factory=dict)

    @model_validator(mode="after")
    def validate_files(self) -> "ZenodoSource":
        for name in self.files:
            safe_remote_path(name)
        if not set(self.published_checksums) <= set(self.files):
            raise ValueError("published checksums must refer to selected Zenodo files")
        return self


class DirectSource(ManifestModel):
    url: str
    filename: str

    @model_validator(mode="after")
    def validate_direct(self) -> "DirectSource":
        if not self.url.startswith(("http://", "https://")):
            raise ValueError("direct URL must use HTTP(S)")
        safe_remote_path(self.filename)
        return self


Source = Union[HuggingFaceSource, GDriveSource, ZenodoSource, DirectSource]


class CheckpointResource(ManifestModel):
    id: str = Field(pattern=r"^[a-z][a-z0-9_]*$")
    backend: Literal["huggingface", "gdrive", "zenodo", "direct"]
    provenance: Provenance
    teacher_family: TeacherFamily
    destination: str
    license: str = Field(min_length=1)
    terms_url: str
    digest: DigestSpec
    required_files: list[str] = Field(min_length=1)
    source: Source
    note: str | None = None

    @model_validator(mode="after")
    def validate_resource(self) -> "CheckpointResource":
        expected = {
            "huggingface": HuggingFaceSource,
            "gdrive": GDriveSource,
            "zenodo": ZenodoSource,
            "direct": DirectSource,
        }[self.backend]
        if not isinstance(self.source, expected):
            raise ValueError(f"{self.backend} resource has the wrong source descriptor")
        path = safe_remote_path(self.destination)
        if path.parts[0] != self.provenance:
            raise ValueError("resource destination must begin with its provenance")
        if not self.terms_url.startswith("https://"):
            raise ValueError("resource terms URL must use HTTPS")
        if len(self.required_files) != len(set(self.required_files)):
            raise ValueError("required checkpoint files must be unique")
        for name in self.required_files:
            safe_remote_path(name)
        if isinstance(self.source, HuggingFaceSource) and any(
            not any(fnmatchcase(name, pattern) for pattern in self.source.allow_patterns)
            for name in self.required_files
        ):
            raise ValueError("required Hugging Face files must be included by allow patterns")
        if isinstance(self.source, GDriveSource) and set(self.required_files) != set(
            self.source.expected_files
        ):
            raise ValueError("Google Drive required files must exactly match pinned file IDs")
        if isinstance(self.source, ZenodoSource) and set(self.required_files) != set(
            self.source.files
        ):
            raise ValueError("Zenodo required files must exactly match selected record files")
        if isinstance(self.source, DirectSource) and self.required_files != [
            self.source.filename
        ]:
            raise ValueError("direct required file must match the declared filename")
        if self.digest.sha256 is not None and len(self.required_files) != 1:
            raise ValueError("aggregate SHA-256 is incompatible with a multi-file resource")
        if self.digest.per_file_sha256 and set(self.digest.per_file_sha256) != set(
            self.required_files
        ):
            raise ValueError("per-file SHA-256 inventory must match required files exactly")
        return self


class CoverageEntry(ManifestModel):
    model: str = Field(pattern=r"^[a-z][a-z0-9_]*$")
    dataset: Literal["lm1b", "owt"]
    resource: str = Field(pattern=r"^[a-z][a-z0-9_]*$")
    path: str | None = None
    teacher_family: TeacherFamily | None = None
    sampling_config: str | None = None
    sampling_config_source: Literal["resource", "project"] | None = None
    sampling_config_sha256: str | None = Field(
        default=None, pattern=r"^[0-9a-f]{64}$"
    )
    sampling_config_source_commit: str | None = Field(
        default=None, pattern=r"^[0-9a-f]{40}$"
    )

    @model_validator(mode="after")
    def validate_path(self) -> "CoverageEntry":
        if self.path is not None:
            safe_remote_path(self.path)
        config_fields = (self.sampling_config, self.sampling_config_source)
        if any(value is None for value in config_fields) != all(
            value is None for value in config_fields
        ):
            raise ValueError("coverage sampling config requires a path and source")
        if self.sampling_config is not None:
            config = safe_remote_path(self.sampling_config)
            if config.suffix.lower() not in {".json", ".yaml", ".yml"}:
                raise ValueError("coverage sampling config has an unsupported suffix")
        if self.sampling_config_source == "project":
            if (
                self.sampling_config_sha256 is None
                or self.sampling_config_source_commit is None
            ):
                raise ValueError(
                    "project sampling config requires a pinned SHA-256 and source commit"
                )
        elif (
            self.sampling_config_sha256 is not None
            or self.sampling_config_source_commit is not None
        ):
            raise ValueError(
                "only project sampling configs declare SHA-256 and source commit"
            )
        return self


class TrainingRecipe(ManifestModel):
    model: str = Field(pattern=r"^[a-z][a-z0-9_]*$")
    dataset: Literal["lm1b", "owt"]
    provenance: Provenance
    teacher_family: TeacherFamily
    source: str = Field(pattern=r"^[a-z][a-z0-9_]*$")
    source_commit: str = Field(pattern=r"^[0-9a-f]{40}$")
    teacher_adapter: TeacherAdapter | None = None
    command: str = Field(min_length=1)
    output: str
    sampling_checkpoint: str | None = None
    sampling_config: str | None = None
    prerequisites: list[str] = Field(default_factory=list)
    note: str | None = None

    @model_validator(mode="after")
    def validate_output(self) -> "TrainingRecipe":
        path = safe_remote_path(self.output)
        if path.parts[0] != "checkpoints" or path.parts[1] != self.provenance:
            raise ValueError("recipe output must be inside checkpoints/<provenance>")
        if self.sampling_checkpoint is not None:
            checkpoint = safe_remote_path(self.sampling_checkpoint)
            if checkpoint.suffix.lower() not in {
                ".bin",
                ".ckpt",
                ".pt",
                ".pth",
                ".safetensors",
            }:
                raise ValueError("recipe sampling checkpoint has an unsupported suffix")
        if self.sampling_config is not None:
            config = safe_remote_path(self.sampling_config)
            if config.suffix.lower() not in {".json", ".yaml", ".yml"}:
                raise ValueError("recipe sampling config has an unsupported suffix")
        if self.source in {"sdtt", "di4c"}:
            expected = {
                "uniform_duo": "uniform_to_absorbing",
                "masked_mdlm": "masked_to_absorbing",
            }.get(self.teacher_family)
            if self.teacher_adapter != expected:
                raise ValueError(
                    f"{self.source} recipe for {self.teacher_family} requires {expected}"
                )
            if self.sampling_checkpoint is None or self.sampling_config is None:
                raise ValueError(
                    f"{self.source} recipe requires exact sampling checkpoint and config"
                )
        elif self.teacher_adapter is not None:
            raise ValueError("teacher adapters are only valid for SDTT/Di4C recipes")
        return self


class CheckpointManifest(ManifestModel):
    model_config = ConfigDict(extra="forbid", populate_by_name=True)

    schema_version: Literal[1]
    resources: dict[str, CheckpointResource]
    recipes: dict[str, TrainingRecipe]
    coverage_entries: list[CoverageEntry] = Field(alias="coverage")

    @field_validator("schema_version", mode="before")
    @classmethod
    def validate_schema_version(cls, value: object) -> object:
        if type(value) is not int or value != 1:
            raise ValueError("checkpoint manifest schema_version must be integer 1")
        return value

    @property
    def coverage(self) -> dict[tuple[str, str], CoverageEntry]:
        return {(item.model, item.dataset): item for item in self.coverage_entries}

    @model_validator(mode="after")
    def validate_links(self) -> "CheckpointManifest":
        for key, resource in self.resources.items():
            if key != resource.id:
                raise ValueError(f"resource key {key} does not match id {resource.id}")
        cells = [(entry.model, entry.dataset) for entry in self.coverage_entries]
        if len(cells) != len(set(cells)):
            raise ValueError("checkpoint coverage cells must be unique")
        for entry in self.coverage_entries:
            if entry.resource not in self.resources:
                raise ValueError(f"unknown checkpoint resource {entry.resource}")
            resource = self.resources[entry.resource]
            if resource.teacher_family == "mixed_baselines":
                if entry.path is None or entry.teacher_family is None:
                    raise ValueError("mixed checkpoint resources require a path and teacher family")
            elif entry.teacher_family not in {None, resource.teacher_family}:
                raise ValueError("coverage teacher family differs from its checkpoint resource")
        return self


def load_checkpoint_manifest(path: Path) -> CheckpointManifest:
    document = yaml.safe_load(path.read_text(encoding="utf-8"))
    if not isinstance(document, dict):
        raise ValueError("checkpoint manifest must be a mapping")
    return CheckpointManifest.model_validate(document)


def validate_checkpoint_coverage(registry, manifest: CheckpointManifest) -> None:
    """Cross-check registry support cells against resources and recipe fallbacks."""

    referenced_recipes: set[str] = set()
    for model_name, model in registry.models.items():
        for dataset, support in model.datasets.items():
            cell = (model_name, dataset)
            coverage = manifest.coverage.get(cell)
            recipe_id = support.train_recipe
            if support.status == "unsupported":
                if coverage is not None:
                    raise ValueError(f"unsupported cell {model_name}/{dataset} has a checkpoint")
                continue
            if (coverage is None) == (recipe_id is None):
                raise ValueError(
                    f"supported cell {model_name}/{dataset} requires exactly one resource or recipe"
                )
            if coverage is not None:
                resource = manifest.resources[coverage.resource]
                if resource.provenance != support.provenance:
                    raise ValueError(f"checkpoint provenance differs for {model_name}/{dataset}")
            else:
                if recipe_id not in manifest.recipes:
                    raise ValueError(f"unknown recipe {recipe_id} for {model_name}/{dataset}")
                recipe = manifest.recipes[recipe_id]
                if (recipe.model, recipe.dataset) != cell:
                    raise ValueError(f"recipe {recipe_id} does not describe {model_name}/{dataset}")
                if recipe.provenance != support.provenance:
                    raise ValueError(f"recipe provenance differs for {model_name}/{dataset}")
                referenced_recipes.add(recipe_id)
    known_cells = {
        (model_name, dataset)
        for model_name, model in registry.models.items()
        for dataset in model.datasets
    }
    extras = set(manifest.coverage) - known_cells
    if extras:
        raise ValueError(f"checkpoint coverage contains unknown cells: {sorted(extras)}")
    if referenced_recipes != set(manifest.recipes):
        raise ValueError("checkpoint manifest contains unreferenced training recipes")


def safe_remote_path(name: str) -> PurePosixPath:
    """Return a safe relative POSIX path controlled by a remote provider."""

    if not name or name in {".", ".."} or "\\" in name:
        raise ValueError(f"unsafe remote path: {name!r}")
    path = PurePosixPath(name)
    if path.is_absolute() or any(part in {"", ".", ".."} for part in path.parts):
        raise ValueError(f"unsafe remote path: {name!r}")
    return path


def require_server_platform(system_name: str) -> None:
    """Permit real checkpoint acquisition only on the approved Linux server."""

    if system_name != "Linux":
        raise RuntimeError(
            f"checkpoint downloads are Linux server-only; detected {system_name or 'unknown'}"
        )


def safe_checkpoint_destination(checkpoint_root: Path, relative_name: str) -> Path:
    """Resolve a manifest path while rejecting symlink ancestors and escapes."""

    relative = safe_remote_path(relative_name)
    root = checkpoint_root.absolute()
    if root.is_symlink():
        raise ValueError(f"checkpoint root is a symlink: {root}")
    candidate = root.joinpath(*relative.parts)
    current = root
    for part in relative.parts:
        current = current / part
        if current.is_symlink():
            raise ValueError(f"checkpoint destination contains a symlink: {current}")
    try:
        candidate.resolve(strict=False).relative_to(root.resolve(strict=False))
    except ValueError as error:
        raise ValueError(f"checkpoint destination escapes root: {candidate}") from error
    return candidate


def verify_published_file(
    path: Path,
    *,
    size_bytes: int | None,
    checksum: str | None,
) -> None:
    """Validate provider-published size/checksum metadata before publication."""

    if size_bytes is not None and path.stat().st_size != size_bytes:
        raise ValueError(f"published size mismatch for {path}")
    if checksum is None:
        return
    try:
        algorithm, expected = checksum.split(":", 1)
    except ValueError as error:
        raise ValueError(f"invalid published checksum for {path}") from error
    if algorithm == "md5":
        digest = hashlib.md5(usedforsecurity=False)
    elif algorithm == "sha256":
        digest = hashlib.sha256()
    else:
        raise ValueError(f"unsupported published checksum algorithm: {algorithm}")
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    if digest.hexdigest() != expected:
        raise ValueError(f"published checksum mismatch for {path}")


def build_hf_snapshot_kwargs(source: HuggingFaceSource, partial: Path) -> dict[str, object]:
    return {
        "repo_id": source.repo_id,
        "revision": source.revision,
        "allow_patterns": list(source.allow_patterns),
        "local_dir": str(partial),
        "resume_download": True,
    }


def build_gdrive_commands(source: GDriveSource, staging: Path) -> list[list[str]]:
    """Address each Drive object by immutable file ID, never folder-discovered name."""

    object_root = staging / ".objects"
    return [
        [
            "gdown",
            "--continue",
            file_id,
            "--output",
            str(object_root / f"{file_id}.partial"),
        ]
        for file_id in source.expected_files.values()
    ]


def _quarantine_path(target: Path, quarantine_root: Path, timestamp: str) -> Path:
    try:
        relative = target.relative_to(quarantine_root.parent)
    except ValueError:
        relative = Path(target.name)
    destination = quarantine_root / timestamp / relative
    destination.parent.mkdir(parents=True, exist_ok=True)
    return destination


def publish_partial(
    partial: Path,
    target: Path,
    digest: DigestSpec,
    *,
    quarantine_root: Path | None = None,
    timestamp: str | None = None,
) -> dict[str, object]:
    """Validate a partial file, quarantine a mismatch, and atomically publish it."""

    if partial.is_symlink():
        raise ValueError(f"checkpoint partial path is a symlink: {partial}")
    if target.is_symlink():
        raise ValueError(f"checkpoint target path is a symlink: {target}")
    observed = sha256_file(partial)
    if digest.sha256 is not None and observed != digest.sha256:
        raise ValueError(f"download digest mismatch for {target}")
    quarantined = None
    if target.exists():
        existing = sha256_file(target)
        if digest.sha256 is not None and existing == digest.sha256:
            partial.unlink()
            return {
                "status": "reused",
                "sha256": existing,
                "size_bytes": target.stat().st_size,
                "quarantined": None,
            }
        quarantine_root = quarantine_root or target.parents[1] / "quarantine"
        timestamp = timestamp or datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
        quarantine = _quarantine_path(target, quarantine_root, timestamp)
        os.replace(target, quarantine)
        quarantined = quarantine.as_posix()
    target.parent.mkdir(parents=True, exist_ok=True)
    os.replace(partial, target)
    return {
        "status": "downloaded",
        "sha256": observed,
        "size_bytes": target.stat().st_size,
        "quarantined": quarantined,
    }


def download_direct(
    source: DirectSource,
    target: Path,
    digest: DigestSpec,
    *,
    quarantine_root: Path | None = None,
    published_size_bytes: int | None = None,
    published_checksum: str | None = None,
) -> dict[str, object]:
    """Resume an HTTP download into ``.partial`` and atomically publish it."""

    quarantined_existing = None
    try:
        target_info = target.lstat()
    except FileNotFoundError:
        target_info = None
    if target_info is not None:
        valid_target = stat.S_ISREG(target_info.st_mode) and target_info.st_size > 0
        if not valid_target:
            effective_quarantine_root = (
                quarantine_root or target.parents[1] / "quarantine"
            )
            stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ")
            quarantine = _quarantine_path(
                target, effective_quarantine_root, stamp
            )
            os.replace(target, quarantine)
            quarantined_existing = quarantine.as_posix()
        elif digest.sha256 and sha256_file(target) == digest.sha256:
            return {
                "status": "reused",
                "sha256": digest.sha256,
                "size_bytes": target_info.st_size,
                "quarantined": None,
            }
    target.parent.mkdir(parents=True, exist_ok=True)
    partial = target.with_name(target.name + ".partial")
    if partial.is_symlink():
        raise ValueError(f"checkpoint partial path is a symlink: {partial}")
    if partial.exists() and not partial.is_file():
        raise ValueError(f"checkpoint partial path is not a file: {partial}")
    metadata_path = partial.with_name(partial.name + ".meta.json")

    def discard_partial() -> None:
        partial.unlink(missing_ok=True)
        metadata_path.unlink(missing_ok=True)

    def load_resume_metadata() -> dict[str, object] | None:
        if not partial.exists() or partial.stat().st_size == 0:
            metadata_path.unlink(missing_ok=True)
            return None
        try:
            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
            if metadata.get("url") != source.url:
                raise ValueError("partial URL differs")
            if metadata.get("validator_header") not in {"ETag", "Last-Modified"}:
                raise ValueError("partial has no supported validator")
            if not isinstance(metadata.get("validator_value"), str) or not metadata[
                "validator_value"
            ]:
                raise ValueError("partial validator is empty")
            total_size = metadata.get("total_size")
            if not isinstance(total_size, int) or total_size < partial.stat().st_size:
                raise ValueError("partial total size is invalid")
            return metadata
        except (OSError, TypeError, ValueError, json.JSONDecodeError):
            discard_partial()
            return None

    def response_header(response, name: str) -> str | None:
        value = response.headers.get(name)
        if value is not None:
            return value
        lowered = name.lower()
        for key, candidate in response.headers.items():
            if key.lower() == lowered:
                return candidate
        return None

    def content_range(response) -> tuple[int, int, int] | None:
        value = response_header(response, "Content-Range")
        if value is None or not value.startswith("bytes "):
            return None
        try:
            span, total = value.removeprefix("bytes ").split("/", 1)
            start, end = span.split("-", 1)
            return int(start), int(end), int(total)
        except (TypeError, ValueError):
            return None

    resume_metadata = load_resume_metadata()
    for attempt in range(2):
        offset = partial.stat().st_size if partial.exists() and resume_metadata else 0
        complete_partial = bool(
            offset and offset == resume_metadata.get("total_size")
        )
        headers = {}
        if offset:
            headers = {
                "Range": "bytes=0-0" if complete_partial else f"bytes={offset}-",
                "If-Range": str(resume_metadata["validator_value"]),
            }
        request = Request(source.url, headers=headers)
        try:
            response_context = urlopen(request)
        except HTTPError as error:
            if error.code == 416 and offset and attempt == 0:
                discard_partial()
                resume_metadata = None
                continue
            raise
        with response_context as response:
            status = getattr(response, "status", response.getcode())
            response_range = content_range(response)
            if offset and status == 416:
                discard_partial()
                resume_metadata = None
                if attempt == 0:
                    continue
                raise ValueError("HTTP range remained unsatisfiable after restart")
            if complete_partial and status == 206:
                validator_header = str(resume_metadata["validator_header"])
                validator = response_header(response, validator_header)
                probe = response.read()
                with partial.open("rb") as handle:
                    first_byte = handle.read(1)
                trusted = (
                    response_range == (0, 0, resume_metadata["total_size"])
                    and validator == resume_metadata["validator_value"]
                    and probe == first_byte
                )
                if trusted:
                    break
                discard_partial()
                resume_metadata = None
                if attempt == 0:
                    continue
                raise ValueError("complete HTTP partial could not be revalidated")
            if offset and status == 206:
                validator_header = str(resume_metadata["validator_header"])
                validator = response_header(response, validator_header)
                trusted = (
                    response_range is not None
                    and response_range[0] == offset
                    and response_range[1] + 1 == response_range[2]
                    and response_range[2] == resume_metadata["total_size"]
                    and validator == resume_metadata["validator_value"]
                )
                if not trusted:
                    discard_partial()
                    resume_metadata = None
                    if attempt == 0:
                        continue
                    raise ValueError("HTTP range response could not be validated")
                mode = "ab"
                total_size = response_range[2]
            elif offset and status == 200:
                mode = "wb"
                length = response_header(response, "Content-Length")
                total_size = int(length) if length is not None else None
            elif not offset and status in {200, 206}:
                if status == 206 and (
                    response_range is None
                    or response_range[0] != 0
                    or response_range[1] + 1 != response_range[2]
                ):
                    raise ValueError("fresh HTTP range response is incomplete")
                mode = "wb"
                length = response_header(response, "Content-Length")
                total_size = (
                    response_range[2]
                    if response_range is not None
                    else int(length) if length is not None else None
                )
            else:
                if offset and attempt == 0:
                    discard_partial()
                    resume_metadata = None
                    continue
                raise ValueError(f"unexpected HTTP download status {status}")

            etag = response_header(response, "ETag")
            modified = response_header(response, "Last-Modified")
            validator_header = "ETag" if etag else "Last-Modified" if modified else None
            validator_value = etag or modified
            if validator_header and validator_value and total_size is not None:
                atomic_json_write(
                    metadata_path,
                    {
                        "url": source.url,
                        "validator_header": validator_header,
                        "validator_value": validator_value,
                        "total_size": total_size,
                    },
                )
            else:
                metadata_path.unlink(missing_ok=True)
            with partial.open(mode) as handle:
                shutil.copyfileobj(response, handle, length=1024 * 1024)
                handle.flush()
                os.fsync(handle.fileno())
            if total_size is not None and partial.stat().st_size != total_size:
                raise ValueError("HTTP response size differs from validated total")
            break
    else:  # pragma: no cover - both attempts either break or raise
        raise RuntimeError("HTTP download retry loop exhausted")
    if partial.stat().st_size == 0:
        raise FileNotFoundError(f"required checkpoint file is empty: {source.filename}")
    verify_published_file(
        partial,
        size_bytes=published_size_bytes,
        checksum=published_checksum,
    )
    metadata_path.unlink(missing_ok=True)
    result = publish_partial(
        partial,
        target,
        digest,
        quarantine_root=quarantine_root,
    )
    if quarantined_existing is not None and result["quarantined"] is None:
        result["quarantined"] = quarantined_existing
    return result


def select_zenodo_files(
    metadata: dict[str, object], source: ZenodoSource
) -> list[dict[str, object]]:
    by_name: dict[str, dict[str, object]] = {}
    for item in metadata.get("files", []):
        name = str(item["key"])
        safe_remote_path(name)
        by_name[name] = item
    selected = []
    for name in source.files:
        safe_remote_path(name)
        if name not in by_name:
            raise FileNotFoundError(f"Zenodo record {source.record_id} lacks {name}")
        item = by_name[name]
        links = item["links"]
        selected.append(
            {
                "filename": name,
                "url": links.get("content") or links["self"],
                "size_bytes": int(item["size"]),
                "published_checksum": item.get("checksum"),
            }
        )
    return selected


def _file_records(path: Path, root: Path) -> list[dict[str, object]]:
    paths = [path] if path.is_file() else sorted(item for item in path.rglob("*") if item.is_file())
    records = []
    for item in paths:
        if item.is_symlink():
            raise ValueError(f"checkpoint tree contains a symlink: {item}")
        safe_remote_path(item.relative_to(path.parent if path.is_file() else path).as_posix())
        records.append(
            {
                "path": item.relative_to(root).as_posix(),
                "size_bytes": item.stat().st_size,
                "sha256": sha256_file(item),
            }
        )
    return records


def _verify_required_files(staging: Path, required_files: list[str]) -> None:
    for relative_name in required_files:
        expected = staging.joinpath(*safe_remote_path(relative_name).parts)
        if not expected.is_file() or expected.is_symlink():
            raise FileNotFoundError(f"required checkpoint file is missing: {relative_name}")
        if expected.stat().st_size == 0:
            raise FileNotFoundError(f"required checkpoint file is empty: {relative_name}")


def _digest_for_file(digest: DigestSpec, relative_name: str) -> DigestSpec:
    if digest.policy == "capture_after_download":
        return digest
    if digest.sha256 is not None:
        return digest
    return DigestSpec(policy="sha256", sha256=digest.per_file_sha256[relative_name])


def _verify_resource_digest(
    staging: Path, digest: DigestSpec, required_files: list[str]
) -> None:
    if digest.policy == "capture_after_download":
        return
    for relative_name in required_files:
        expected = _digest_for_file(digest, relative_name).sha256
        path = staging.joinpath(*safe_remote_path(relative_name).parts)
        if expected is None or sha256_file(path) != expected:
            raise ValueError(f"download digest mismatch for {path}")


def _quarantine_unexpected_gdrive_staging(
    partial: Path,
    source: GDriveSource,
    quarantine_root: Path,
) -> list[str]:
    allowed_files = set(source.expected_files)
    allowed_object_files = {
        f".objects/{file_id}.partial" for file_id in source.expected_files.values()
    }
    allowed_files.update(allowed_object_files)
    allowed_directories = {".objects"}
    for name in allowed_files:
        relative = PurePosixPath(name)
        allowed_directories.update(
            parent.as_posix() for parent in relative.parents if parent.as_posix() != "."
        )
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ")
    quarantined = []
    for name in allowed_object_files:
        relative = PurePosixPath(name)
        current = partial
        valid_parents = True
        for part in relative.parts[:-1]:
            current = current / part
            try:
                parent_info = current.lstat()
            except FileNotFoundError:
                valid_parents = False
                break
            if not stat.S_ISDIR(parent_info.st_mode):
                valid_parents = False
                break
        if not valid_parents:
            continue
        item = partial.joinpath(*relative.parts)
        try:
            item_info = item.lstat()
        except FileNotFoundError:
            continue
        if stat.S_ISREG(item_info.st_mode) and item_info.st_size > 0:
            continue
        quarantine = _quarantine_path(item, quarantine_root, stamp)
        os.replace(item, quarantine)
        quarantined.append(quarantine.as_posix())
    for item in sorted(
        partial.rglob("*"), key=lambda path: len(path.parts), reverse=True
    ):
        relative_name = item.relative_to(partial).as_posix()
        if item.is_symlink() or not item.is_dir():
            if relative_name not in allowed_files:
                quarantine = _quarantine_path(item, quarantine_root, stamp)
                os.replace(item, quarantine)
                quarantined.append(quarantine.as_posix())
        elif relative_name not in allowed_directories:
            try:
                item.rmdir()
            except OSError:
                pass
    return quarantined


def _publish_directory(partial: Path, destination: Path, quarantine_root: Path) -> str | None:
    if partial.is_symlink():
        raise ValueError(f"checkpoint partial path is a symlink: {partial}")
    for item in partial.rglob("*"):
        if item.is_symlink():
            raise ValueError(f"checkpoint tree contains a symlink: {item}")
        safe_remote_path(item.relative_to(partial).as_posix())
    quarantined = None
    if destination.exists():
        stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
        quarantine = _quarantine_path(destination, quarantine_root, stamp)
        os.replace(destination, quarantine)
        quarantined = quarantine.as_posix()
    destination.parent.mkdir(parents=True, exist_ok=True)
    os.replace(partial, destination)
    return quarantined


def fetch_resource(root: Path, resource: CheckpointResource) -> dict[str, object]:
    """Fetch one resource and return a lock-ready status record."""

    checkpoint_root = root / "checkpoints"
    destination = safe_checkpoint_destination(checkpoint_root, resource.destination)
    quarantine_root = safe_checkpoint_destination(checkpoint_root, "quarantine")
    if resource.backend == "direct":
        source = resource.source
        target = safe_checkpoint_destination(
            checkpoint_root, f"{resource.destination}/{source.filename}"
        )
        result = download_direct(
            source,
            target,
            _digest_for_file(resource.digest, source.filename),
            quarantine_root=quarantine_root,
        )
        quarantined = result.get("quarantined")
    elif resource.backend == "huggingface":
        source = resource.source
        partial = safe_checkpoint_destination(
            checkpoint_root, resource.destination + ".partial"
        )
        if partial.exists() and not partial.is_dir():
            raise ValueError(f"HF partial path is not a directory: {partial}")
        partial.mkdir(parents=True, exist_ok=True)
        from huggingface_hub import snapshot_download

        snapshot_download(**build_hf_snapshot_kwargs(source, partial))
        _verify_required_files(partial, resource.required_files)
        _verify_resource_digest(partial, resource.digest, resource.required_files)
        quarantined = _publish_directory(partial, destination, quarantine_root)
    elif resource.backend == "gdrive":
        source = resource.source
        partial = safe_checkpoint_destination(
            checkpoint_root, resource.destination + ".partial"
        )
        partial.mkdir(parents=True, exist_ok=True)
        _quarantine_unexpected_gdrive_staging(partial, source, quarantine_root)
        object_root = partial / ".objects"
        object_root.mkdir(parents=True, exist_ok=True)
        commands = build_gdrive_commands(source, partial)
        for (relative_name, file_id), command in zip(
            source.expected_files.items(), commands
        ):
            object_path = safe_checkpoint_destination(
                checkpoint_root,
                f"{resource.destination}.partial/.objects/{file_id}.partial",
            )
            subprocess.run(command, check=True)
            if not object_path.is_file() or object_path.is_symlink():
                raise FileNotFoundError(f"Google Drive file ID {file_id} produced no file")
            declared_path = safe_checkpoint_destination(
                checkpoint_root, f"{resource.destination}.partial/{relative_name}"
            )
            declared_path.parent.mkdir(parents=True, exist_ok=True)
            os.replace(object_path, declared_path)
        object_root.rmdir()
        _verify_required_files(partial, resource.required_files)
        observed_inventory = {
            item.relative_to(partial).as_posix()
            for item in partial.rglob("*")
            if item.is_file()
        }
        if observed_inventory != set(source.expected_files):
            raise ValueError("Google Drive staging inventory differs from pinned ID/path map")
        _verify_resource_digest(partial, resource.digest, resource.required_files)
        quarantined = _publish_directory(partial, destination, quarantine_root)
    elif resource.backend == "zenodo":
        source = resource.source
        api_url = f"https://zenodo.org/api/records/{source.record_id}"
        with urlopen(api_url) as response:
            metadata = json.load(response)
        destination.mkdir(parents=True, exist_ok=True)
        quarantined = None
        for item in select_zenodo_files(metadata, source):
            filename = str(item["filename"])
            published_checksum = item.get("published_checksum")
            pinned_checksum = source.published_checksums.get(filename)
            if pinned_checksum is not None and published_checksum != pinned_checksum:
                raise ValueError(f"Zenodo metadata checksum drift for {filename}")
            result = download_direct(
                DirectSource(url=str(item["url"]), filename=filename),
                safe_checkpoint_destination(
                    checkpoint_root, f"{resource.destination}/{filename}"
                ),
                _digest_for_file(resource.digest, filename),
                quarantine_root=quarantine_root,
                published_size_bytes=int(item["size_bytes"]),
                published_checksum=pinned_checksum or published_checksum,
            )
            quarantined = quarantined or result.get("quarantined")
        _verify_required_files(destination, resource.required_files)
    else:  # pragma: no cover - Pydantic makes this unreachable
        raise ValueError(f"unsupported backend {resource.backend}")
    files = _file_records(destination, root)
    if not files:
        raise FileNotFoundError(f"checkpoint resource {resource.id} produced no files")
    return {
        "status": "downloaded",
        "backend": resource.backend,
        "provenance": resource.provenance,
        "teacher_family": resource.teacher_family,
        "destination": destination.relative_to(root).as_posix(),
        "quarantined": quarantined,
        "files": files,
    }


def fetch_all_resources(root: Path, manifest_path: Path, manifest: CheckpointManifest) -> dict[str, object]:
    records: dict[str, object] = {}
    for resource_id, resource in manifest.resources.items():
        try:
            records[resource_id] = fetch_resource(root, resource)
        except Exception as error:  # preserve provider failures as structured lock evidence
            status = "unavailable" if isinstance(error, (FileNotFoundError, HTTPError)) else "error"
            records[resource_id] = {
                "status": status,
                "backend": resource.backend,
                "provenance": resource.provenance,
                "teacher_family": resource.teacher_family,
                "destination": f"checkpoints/{resource.destination}",
                "files": [],
                "error": {"type": type(error).__name__, "message": str(error)},
            }
    lock = {
        "schema_version": 1,
        "created_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "manifest_sha256": sha256_file(manifest_path),
        "resources": records,
    }
    atomic_json_write(root / "artifacts" / "checkpoint_lock.json", lock)
    return lock


def _lstat_regular_checkpoint_path(root: Path, relative: PurePosixPath) -> Path:
    """Reject symlinks/non-directories at every component and require a regular file."""

    current = root.absolute()
    for index, part in enumerate(relative.parts):
        current = current / part
        info = current.lstat()
        if stat.S_ISLNK(info.st_mode):
            raise ValueError(f"checkpoint lock path contains a symlink: {current}")
        final = index == len(relative.parts) - 1
        if final and not stat.S_ISREG(info.st_mode):
            raise ValueError(f"checkpoint lock path is not a regular file: {current}")
        if not final and not stat.S_ISDIR(info.st_mode):
            raise ValueError(f"checkpoint lock ancestor is not a directory: {current}")
    return current


def verify_checkpoint_lock(
    root: Path,
    lock: dict[str, object],
    *,
    manifest_path: Path | None = None,
) -> dict[str, object]:
    report: dict[str, object] = {
        "ok": True,
        "schema_status": (
            "verified"
            if type(lock.get("schema_version")) is int and lock.get("schema_version") == 1
            else "invalid"
        ),
        "resources": {},
    }
    if report["schema_status"] != "verified":
        report["ok"] = False
    resources = lock.get("resources")
    if not isinstance(resources, dict) or not resources:
        report["resource_set_status"] = "invalid"
        report["ok"] = False
        resources = resources if isinstance(resources, dict) else {}
    else:
        report["resource_set_status"] = "verified"
    manifest = None
    if manifest_path is not None:
        expected_manifest = lock.get("manifest_sha256")
        observed_manifest = sha256_file(manifest_path) if manifest_path.is_file() else None
        report["manifest_status"] = (
            "verified" if observed_manifest == expected_manifest else "mismatch"
        )
        if report["manifest_status"] != "verified":
            report["ok"] = False
        try:
            manifest = load_checkpoint_manifest(manifest_path)
            expected_ids = set(manifest.resources)
        except (OSError, ValueError):
            report["resource_set_status"] = "invalid"
            report["ok"] = False
        else:
            if set(resources) != expected_ids:
                report["resource_set_status"] = "mismatch"
                report["ok"] = False
    for resource_id, record in resources.items():
        if not isinstance(record, dict):
            report["resources"][resource_id] = {
                "status": "invalid",
                "files": [],
                "error": {
                    "type": "InvalidResourceRecord",
                    "message": f"lock resource {resource_id} must be a mapping",
                },
            }
            report["ok"] = False
            continue
        manifest_resource = manifest.resources.get(resource_id) if manifest is not None else None
        if manifest_resource is not None:
            expected_metadata = {
                "backend": manifest_resource.backend,
                "provenance": manifest_resource.provenance,
                "teacher_family": manifest_resource.teacher_family,
                "destination": f"checkpoints/{manifest_resource.destination}",
            }
            if any(record.get(field) != value for field, value in expected_metadata.items()):
                report["resources"][resource_id] = {
                    "status": "invalid",
                    "error": {"message": "lock resource metadata differs from manifest"},
                }
                report["ok"] = False
                continue
        if record.get("status") not in {"downloaded", "verified"}:
            report["resources"][resource_id] = {
                "status": record.get("status", "unavailable"),
                "error": record.get("error"),
            }
            report["ok"] = False
            continue
        files = record.get("files")
        if not isinstance(files, list) or not files:
            report["resources"][resource_id] = {"status": "invalid", "files": []}
            report["ok"] = False
            continue
        status = "verified"
        details = []
        seen_paths: set[str] = set()
        relative_to_destination: set[str] = set()
        destination_prefix = PurePosixPath(
            f"checkpoints/{manifest_resource.destination}"
            if manifest_resource is not None
            else "checkpoints"
        )
        for expected in files:
            try:
                if not isinstance(expected, dict):
                    raise ValueError("lock file record must be a mapping")
                relative = safe_remote_path(str(expected["path"]))
                if relative == destination_prefix or not relative.is_relative_to(
                    destination_prefix
                ):
                    raise ValueError("lock file lies outside its resource destination")
                path_text = relative.as_posix()
                if path_text in seen_paths:
                    raise ValueError("lock contains a duplicate file path")
                seen_paths.add(path_text)
                size_bytes = expected["size_bytes"]
                digest = expected["sha256"]
                if not isinstance(size_bytes, int) or size_bytes < 0:
                    raise ValueError("invalid locked file size")
                if not isinstance(digest, str) or len(digest) != 64 or any(
                    character not in "0123456789abcdef" for character in digest
                ):
                    raise ValueError("invalid locked SHA-256")
                relative_to_destination.add(
                    relative.relative_to(destination_prefix).as_posix()
                )
                path = _lstat_regular_checkpoint_path(root, relative)
            except FileNotFoundError:
                status = "missing"
                details.append({"path": expected.get("path"), "status": "missing"})
                continue
            except (KeyError, TypeError, ValueError, OSError):
                status = "invalid"
                details.append({"path": expected.get("path"), "status": "invalid"})
                break
            if path.stat().st_size != size_bytes or sha256_file(path) != digest:
                status = "mismatch"
                details.append({"path": expected["path"], "status": "mismatch"})
            else:
                details.append({"path": expected["path"], "status": "verified"})
        if (
            status == "verified"
            and manifest_resource is not None
            and not set(manifest_resource.required_files) <= relative_to_destination
        ):
            status = "invalid"
            details.append({"status": "invalid", "error": "required file missing from lock"})
        report["resources"][resource_id] = {"status": status, "files": details}
        if status != "verified":
            report["ok"] = False
    return report