scripts/build_upload_folders.py

#!/usr/bin/env python3
"""Build curated, integrity-checked Kaggle and GitHub upload folders."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import sys
import tempfile
from collections import defaultdict
from pathlib import Path
from urllib.parse import unquote, urlsplit

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from scripts.publish_hf import (  # noqa: E402
    _validate_inventory,
    _validate_provenance,
    _validated_release_status,
)

_KAGGLE_TEMPLATE_FILES = (
    "README.md",
    "example_inference.py",
    "requirements.txt",
    "smoke_test.py",
    "turn_detector.py",
)
_ROOT_GITHUB_FILES = (
    ".env.example",
    ".gitignore",
    ".python-version",
    "DATA_CARD.md",
    "LICENSE",
    "MODEL_CARD.md",
    "Makefile",
    "NOTICE",
    "README.md",
    "REPORT.md",
    "app.py",
    "pyproject.toml",
    "requirements-export.txt",
    "requirements-publish.txt",
    "requirements-training.txt",
)
_GITHUB_ARTIFACTS = (
    "artifacts/partial-baseline/model.json",
    "artifacts/partial-shard-preview-tinytcn-4s/best.pt",
    "artifacts/partial-shard-warmstart-lr3e4-5ep/best.pt",
    "artifacts/partial-shard-warmstart-lr3e4-5ep/cpu_benchmark.json",
    "artifacts/partial-shard-warmstart-lr3e4-5ep/export_manifest.json",
    "artifacts/partial-shard-warmstart-lr3e4-5ep/history.json",
    "artifacts/partial-shard-warmstart-lr3e4-5ep/model.onnx",
    "artifacts/partial-shard-warmstart-lr3e4-5ep/model_metadata.json",
    "artifacts/partial-shard-warmstart-lr3e4-5ep/onnx_benchmark.json",
    "artifacts/partial-shard-warmstart-lr3e4-5ep/onnx_e2e_benchmark.json",
    "artifacts/partial-shard-warmstart-lr3e4-5ep/resolved_config.json",
)
_COLLECTION_FILES = (
    "assignments.jsonl",
    "assignments.summary.json",
    "controller_replay_fixture.jsonl",
)
_PUBLIC_REPORTS = (
    "controller_replay_integration.jsonl",
    "controller_replay_integration.summary.json",
    "dataset_snapshot_status.json",
    "environment_snapshot.json",
    "partial_baseline_metrics.json",
    "partial_iid_split.json",
    "partial_shard_audit.json",
    "partial_source_holdout_baseline_metrics.json",
    "partial_source_holdout_comparison.json",
    "partial_source_holdout_split.json",
    "partial_source_holdout_tinytcn_metrics.json",
    "partial_tinytcn_comparison.json",
    "partial_tinytcn_e2e_benchmark.json",
    "partial_tinytcn_failures.json",
    "partial_tinytcn_metrics.json",
    "partial_tinytcn_onnx_benchmark.json",
    "partial_tinytcn_pytorch_benchmark.json",
    "partial_tinytcn_silence_sensitivity.json",
)
_SECRET_PATTERNS = {
    "Hugging Face token": re.compile(r"\bhf_[A-Za-z0-9]{20,}\b"),
    "GitHub token": re.compile(r"\b(?:ghp|github_pat)_[A-Za-z0-9_]{20,}\b"),
    "OpenAI-style token": re.compile(r"\bsk-[A-Za-z0-9_-]{20,}\b"),
    "private key": re.compile(r"-----BEGIN (?:RSA |EC |OPENSSH )?PRIVATE KEY-----"),
    "absolute macOS user path": re.compile("/" + r"Users/[^/\s]+/"),
}
_MARKDOWN_LINK = re.compile(r"!?\[[^\]]*\]\(([^)]+)\)")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--release-dir", default="release")
    parser.add_argument("--output", default="upload-ready")
    parser.add_argument(
        "--allow-development-release",
        action="store_true",
        help="package the visibly labelled development preview",
    )
    return parser.parse_args()


def _resolve_inside_root(value: str, *, purpose: str) -> Path:
    unresolved = Path(value)
    path = (unresolved if unresolved.is_absolute() else ROOT / unresolved).resolve()
    try:
        path.relative_to(ROOT)
    except ValueError as exc:
        raise SystemExit(f"{purpose} must stay inside the project directory") from exc
    return path


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


def _copy_file(source: Path, destination: Path) -> None:
    if source.is_symlink() or not source.is_file():
        raise SystemExit(f"refusing non-regular upload source: {source}")
    destination.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(source, destination)


def _copy_tree(source: Path, destination: Path, *, suffixes: set[str] | None = None) -> None:
    if source.is_symlink() or not source.is_dir():
        raise SystemExit(f"upload source directory is missing or unsafe: {source}")
    for path in sorted(source.rglob("*")):
        if path.is_symlink():
            raise SystemExit(f"refusing symlink in upload source: {path}")
        if not path.is_file():
            continue
        if "__pycache__" in path.parts or path.suffix == ".pyc":
            continue
        if suffixes is not None and path.suffix not in suffixes:
            continue
        _copy_file(path, destination / path.relative_to(source))


def _load_validated_release(release_dir: Path) -> tuple[dict, bool]:
    manifest_path = release_dir / "release_manifest.json"
    try:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise SystemExit("release_manifest.json is missing or invalid; rebuild release") from exc
    if not isinstance(manifest, dict):
        raise SystemExit("release_manifest.json must contain an object")
    _validate_inventory(release_dir, manifest)
    _validate_provenance(release_dir)
    development_only = _validated_release_status(release_dir, manifest)
    return manifest, development_only


def _require_identical(left: Path, right: Path, label: str) -> None:
    if left.stat().st_size != right.stat().st_size or _sha256(left) != _sha256(right):
        raise SystemExit(f"{label} differs from the validated Hugging Face release")


def _build_kaggle_folder(release_dir: Path, destination: Path, metrics_name: str) -> None:
    model_release = release_dir / "model"
    for name in _KAGGLE_TEMPLATE_FILES:
        _copy_file(ROOT / "deployment" / "kaggle" / name, destination / name)
    for source_name, destination_name in (
        ("model.onnx", "model.onnx"),
        ("model_metadata.json", "model_metadata.json"),
        ("export_manifest.json", "export_manifest.json"),
        (
            metrics_name,
            "development_metrics.json" if "development" in metrics_name else "test_metrics.json",
        ),
        ("MODEL_CARD.md", "MODEL_CARD.md"),
        ("DATA_CARD.md", "DATA_CARD.md"),
        ("LICENSE", "LICENSE"),
        ("NOTICE", "NOTICE"),
    ):
        _copy_file(model_release / source_name, destination / destination_name)
    _copy_file(
        model_release / "run_artifacts" / "onnx_e2e_benchmark.json",
        destination / "benchmark.json",
    )


def _github_gitignore() -> str:
    base = (ROOT / ".gitignore").read_text(encoding="utf-8").rstrip()
    by_parent: dict[Path, list[Path]] = defaultdict(list)
    for relative in map(Path, _GITHUB_ARTIFACTS):
        by_parent[relative.parent].append(relative)
    lines = [
        base,
        "",
        "# Exact public-preview artifacts included by build_upload_folders.py",
        "!artifacts/",
        "artifacts/*",
    ]
    for parent in sorted(by_parent, key=lambda item: item.as_posix()):
        lines.append(f"!{parent.as_posix()}/")
        lines.append(f"{parent.as_posix()}/*")
        for relative in sorted(by_parent[parent], key=lambda item: item.as_posix()):
            lines.append(f"!{relative.as_posix()}")
    return "\n".join(lines) + "\n"


def _build_github_folder(destination: Path) -> None:
    for name in _ROOT_GITHUB_FILES:
        _copy_file(ROOT / name, destination / name)
    (destination / ".gitignore").write_text(_github_gitignore(), encoding="utf-8")

    _copy_tree(ROOT / ".github", destination / ".github")
    _copy_tree(ROOT / "src", destination / "src", suffixes={".py", ".typed"})
    _copy_tree(ROOT / "scripts", destination / "scripts", suffixes={".py", ".sh"})
    _copy_tree(ROOT / "configs", destination / "configs", suffixes={".json", ".yaml", ".yml"})
    _copy_tree(ROOT / "tests", destination / "tests", suffixes={".py"})
    _copy_tree(ROOT / "docs", destination / "docs", suffixes={".md"})
    _copy_tree(ROOT / "deployment", destination / "deployment", suffixes={".md", ".py", ".txt"})

    for name in ("README.md", "requirements.txt"):
        _copy_file(ROOT / "space" / name, destination / "space" / name)
    for name in _COLLECTION_FILES:
        _copy_file(
            ROOT / "data" / "collection" / name,
            destination / "data" / "collection" / name,
        )
    for name in _PUBLIC_REPORTS:
        _copy_file(ROOT / "reports" / name, destination / "reports" / name)
    for relative in _GITHUB_ARTIFACTS:
        _copy_file(ROOT / relative, destination / relative)
    _copy_file(
        ROOT / "artifacts/partial-baseline/model.json",
        destination / "reference_models/acoustic_baseline.json",
    )


def _inventory(
    folder: Path, *, excluded: set[str] | None = None
) -> dict[str, dict[str, int | str]]:
    excluded = excluded or set()
    return {
        path.relative_to(folder).as_posix(): {
            "bytes": path.stat().st_size,
            "sha256": _sha256(path),
        }
        for path in sorted(folder.rglob("*"))
        if path.is_file() and path.relative_to(folder).as_posix() not in excluded
    }


def _write_package_evidence(
    folder: Path,
    *,
    kind: str,
    development_only: bool,
    release_inventory_sha256: str,
) -> None:
    manifest_path = folder / "UPLOAD_MANIFEST.json"
    checksums_path = folder / "SHA256SUMS"
    payload = {
        "format_version": 1,
        "package_kind": kind,
        "development_only": development_only,
        "source_release_inventory_sha256": release_inventory_sha256,
        "files": _inventory(folder),
    }
    manifest_path.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    checksums = _inventory(folder, excluded={checksums_path.name})
    checksums_path.write_text(
        "".join(f"{evidence['sha256']}  {relative}\n" for relative, evidence in checksums.items()),
        encoding="utf-8",
    )


def _secret_findings(folder: Path) -> list[str]:
    findings: list[str] = []
    for path in sorted(folder.rglob("*")):
        if path.is_symlink():
            findings.append(f"symlink: {path.relative_to(folder)}")
            continue
        if not path.is_file() or path.suffix in {".onnx", ".pt"}:
            continue
        if path.name == ".env" or path.name == "hf_publish_receipt.json":
            findings.append(f"forbidden file: {path.relative_to(folder)}")
            continue
        try:
            text = path.read_text(encoding="utf-8")
        except UnicodeDecodeError:
            continue
        for label, pattern in _SECRET_PATTERNS.items():
            if pattern.search(text):
                findings.append(f"{label}: {path.relative_to(folder)}")
    return findings


def _broken_local_markdown_links(folder: Path) -> list[str]:
    broken: list[str] = []
    for markdown in sorted(folder.rglob("*.md")):
        text = markdown.read_text(encoding="utf-8")
        for match in _MARKDOWN_LINK.finditer(text):
            raw = match.group(1).strip().strip("<>")
            raw = raw.split(maxsplit=1)[0].strip("\"'")
            parsed = urlsplit(raw)
            if parsed.scheme or parsed.netloc or not parsed.path:
                continue
            target = (markdown.parent / unquote(parsed.path)).resolve()
            try:
                target.relative_to(folder.resolve())
            except ValueError:
                broken.append(f"{markdown.relative_to(folder)} -> {raw}")
                continue
            if not target.exists():
                broken.append(f"{markdown.relative_to(folder)} -> {raw}")
    return broken


def _verify_sha256sums(folder: Path) -> None:
    path = folder / "SHA256SUMS"
    for line in path.read_text(encoding="utf-8").splitlines():
        expected, relative = line.split("  ", 1)
        target = folder / relative
        if not target.is_file() or _sha256(target) != expected:
            raise SystemExit(f"generated checksum failed validation: {folder.name}/{relative}")


def _validate_output(kaggle: Path, github: Path) -> None:
    nested_kaggle = [path for path in kaggle.rglob("*") if path.is_file() and path.parent != kaggle]
    if nested_kaggle:
        raise SystemExit(f"Kaggle drag-and-drop package is not flat: {nested_kaggle[0]}")
    for folder in (kaggle, github):
        findings = _secret_findings(folder)
        if findings:
            raise SystemExit(f"unsafe generated upload package: {findings[0]}")
        _verify_sha256sums(folder)
    broken = _broken_local_markdown_links(github)
    if broken:
        raise SystemExit(f"generated GitHub package has a broken local link: {broken[0]}")
    oversized = [
        path for path in github.rglob("*") if path.is_file() and path.stat().st_size >= 100_000_000
    ]
    if oversized:
        raise SystemExit(f"generated GitHub file exceeds 100 MB: {oversized[0]}")


def _replace_directory(staged: Path, destination: Path) -> None:
    backup: Path | None = None
    if destination.exists():
        backup = Path(
            tempfile.mkdtemp(prefix=f".{destination.name}.previous-", dir=destination.parent)
        )
        backup.rmdir()
        os.replace(destination, backup)
    try:
        os.replace(staged, destination)
    except Exception:
        if backup is not None and backup.exists() and not destination.exists():
            os.replace(backup, destination)
        raise
    if backup is not None:
        shutil.rmtree(backup)


def main() -> int:
    args = parse_args()
    release_dir = _resolve_inside_root(args.release_dir, purpose="--release-dir")
    destination = _resolve_inside_root(args.output, purpose="--output")
    relative_output = destination.relative_to(ROOT)
    valid_output_root = bool(relative_output.parts) and (
        relative_output.parts[0] == "upload-ready"
        or relative_output.parts[0].startswith("upload-ready-")
    )
    if not valid_output_root:
        raise SystemExit("--output must be upload-ready/ or upload-ready-*/")
    if (
        destination == ROOT
        or release_dir == destination
        or destination in release_dir.parents
        or release_dir in destination.parents
    ):
        raise SystemExit("upload output overlaps the project or release input")
    if Path(args.output).is_symlink():
        raise SystemExit("refusing a symlink upload destination")

    manifest, development_only = _load_validated_release(release_dir)
    if development_only and not args.allow_development_release:
        raise SystemExit("development preview requires --allow-development-release")
    model_release = release_dir / "model"
    _require_identical(
        ROOT / "artifacts/partial-shard-warmstart-lr3e4-5ep/model.onnx",
        model_release / "model.onnx",
        "ONNX artifact",
    )
    _require_identical(
        ROOT / "artifacts/partial-shard-warmstart-lr3e4-5ep/model_metadata.json",
        model_release / "model_metadata.json",
        "model metadata",
    )
    _require_identical(
        ROOT / "artifacts/partial-shard-warmstart-lr3e4-5ep/export_manifest.json",
        model_release / "export_manifest.json",
        "export manifest",
    )

    destination.parent.mkdir(parents=True, exist_ok=True)
    staged = Path(tempfile.mkdtemp(prefix=f".{destination.name}.staging-", dir=destination.parent))
    kaggle = staged / "kaggle-model"
    github = staged / "github-repository"
    kaggle.mkdir()
    github.mkdir()
    metrics_name = "development_metrics.json" if development_only else "test_metrics.json"
    _build_kaggle_folder(release_dir, kaggle, metrics_name)
    _build_github_folder(github)
    release_digest = str(manifest["release_inventory_sha256"])
    _write_package_evidence(
        kaggle,
        kind="kaggle_model_drag_and_drop",
        development_only=development_only,
        release_inventory_sha256=release_digest,
    )
    _write_package_evidence(
        github,
        kind="github_repository",
        development_only=development_only,
        release_inventory_sha256=release_digest,
    )
    _validate_output(kaggle, github)
    summary = {
        "format_version": 1,
        "development_only": development_only,
        "source_release_inventory_sha256": release_digest,
        "folders": {
            "kaggle-model": {
                "files": len(_inventory(kaggle)),
                "bytes": sum(path.stat().st_size for path in kaggle.iterdir() if path.is_file()),
            },
            "github-repository": {
                "files": len(_inventory(github)),
                "bytes": sum(path.stat().st_size for path in github.rglob("*") if path.is_file()),
            },
        },
    }
    (staged / "README.md").write_text(
        "# Upload-ready folders\n\n"
        "- Upload the **contents** of `kaggle-model/` as one Kaggle ONNX model variation.\n"
        "- Upload the **contents** of `github-repository/` to a new GitHub repository.\n"
        "- Both packages are development previews; read their cards before changing visibility.\n",
        encoding="utf-8",
    )
    (staged / "UPLOAD_SUMMARY.json").write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    _replace_directory(staged, destination)
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


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