source/src/bgc_retrieval/artifacts.py

"""Content-addressed manifests and immutable run directories."""

from __future__ import annotations

import hashlib
import json
import os
import platform
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable


def sha256_file(path: str | Path, chunk_size: int = 8 * 1024 * 1024) -> str:
    digest = hashlib.sha256()
    with Path(path).open("rb") as handle:
        while chunk := handle.read(chunk_size):
            digest.update(chunk)
    return digest.hexdigest()


def sha256_json(value: Any) -> str:
    encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def file_record(path: str | Path, root: str | Path | None = None) -> dict[str, Any]:
    file_path = Path(path).resolve()
    stat = file_path.stat()
    display = file_path.relative_to(Path(root).resolve()) if root else file_path
    return {
        "path": str(display),
        "size_bytes": stat.st_size,
        "modified_ns": stat.st_mtime_ns,
        "sha256": sha256_file(file_path),
    }


def environment_record() -> dict[str, Any]:
    packages: dict[str, str] = {}
    for module_name in ("torch", "numpy", "pandas", "scipy", "h5py", "yaml"):
        try:
            module = __import__(module_name)
            packages[module_name] = str(getattr(module, "__version__", "unknown"))
        except ImportError:
            packages[module_name] = "not-installed"
    return {
        "created_at": datetime.now(timezone.utc).isoformat(),
        "python": sys.version,
        "platform": platform.platform(),
        "hostname": platform.node(),
        "packages": packages,
    }


def git_record(repository: str | Path) -> dict[str, Any]:
    repo = Path(repository)

    def run(*args: str) -> str:
        completed = subprocess.run(
            ["git", *args], cwd=repo, text=True, capture_output=True, check=False
        )
        return completed.stdout.strip() if completed.returncode == 0 else "unavailable"

    return {"commit": run("rev-parse", "HEAD"), "status": run("status", "--short")}


def write_json_immutable(path: str | Path, value: Any) -> None:
    output = Path(path)
    output.parent.mkdir(parents=True, exist_ok=True)
    if output.exists():
        raise FileExistsError(f"Refusing to overwrite immutable artifact: {output}")
    temporary = output.with_suffix(output.suffix + f".tmp-{os.getpid()}")
    with temporary.open("w", encoding="utf-8") as handle:
        json.dump(value, handle, indent=2, sort_keys=True)
        handle.write("\n")
    temporary.replace(output)


def create_run_directory(run_root: str | Path, run_id: str) -> Path:
    run_dir = Path(run_root) / run_id
    run_dir.mkdir(parents=True, exist_ok=False)
    return run_dir


def build_file_manifest(paths: Iterable[str | Path], root: str | Path) -> dict[str, Any]:
    root_path = Path(root).resolve()
    records = [file_record(path, root_path) for path in paths if Path(path).is_file()]
    return {
        "schema_version": 1,
        "root": str(root_path),
        "files": records,
        "manifest_sha256": sha256_json(records),
        "environment": environment_record(),
        "git": git_record(root_path),
    }