airfRANS-model-exploration/src/airfrans_frontier/remote/collection.py

from __future__ import annotations

from collections.abc import Callable, Iterable, Mapping
import json
from pathlib import Path
import time
from typing import Any

ARTIFACT_COLLECTION_REPORT = "artifact_collection_report.json"
_PARTIAL_SUFFIXES = (".tmp", ".part", ".partial")
_RSYNC_TEMP_DIRS = (".rsync-partial", ".~tmp~")


def collect_artifact_paths(
    *,
    local_dir: str | Path,
    remote_dir: str | Path,
    paths: Iterable[str],
    required: Iterable[str] = (),
    collection_kind: str,
    copy_one: Callable[[str], int | None],
    clock: Callable[[], float] | None = None,
) -> dict[str, Any]:
    """Copy artifact paths and update artifact_collection_report.json.

    copy_one receives each relative artifact path. It may raise or return a
    non-zero return code; both are recorded per path without losing the rest of
    the collection report.
    """

    now = clock or time.time
    root = Path(local_dir)
    root.mkdir(parents=True, exist_ok=True)
    report_path = root / ARTIFACT_COLLECTION_REPORT
    report = _load_report(report_path)
    required_set = set(required)
    attempted_paths = tuple(dict.fromkeys(paths))
    batch_started_at = now()
    batch_id = f"{collection_kind}-{int(batch_started_at * 1000)}-{len(report['attempts'])}"
    report["batches"].append(
        {
            "batch_id": batch_id,
            "collection_kind": collection_kind,
            "started_at": batch_started_at,
            "paths": list(attempted_paths),
        }
    )
    _write_report(report_path, _refresh_summary(report, now=now()))

    for relative_path in attempted_paths:
        _validate_relative_path(relative_path)
        source = f"{str(remote_dir).rstrip('/')}/{relative_path}"
        destination = root / relative_path
        started = now()
        attempt: dict[str, Any] = {
            "batch_id": batch_id,
            "collection_kind": collection_kind,
            "expected_path": relative_path,
            "required": relative_path in required_set,
            "source_path": source,
            "local_destination": str(destination),
            "attempted": True,
            "started_at": started,
            "bytes_copied": None,
            "duration_seconds": None,
            "return_code": None,
            "exception": None,
            "final_status": "failed",
            "likely_reason": None,
        }
        try:
            return_code = copy_one(relative_path)
            if return_code is not None:
                attempt["return_code"] = int(return_code)
        except Exception as exc:
            attempt["exception"] = {"type": type(exc).__name__, "message": str(exc)}
            attempt["final_status"] = "failed"
            attempt["likely_reason"] = "collection_command_failed"
        else:
            if attempt["return_code"] not in (None, 0):
                attempt["final_status"] = "failed"
                attempt["likely_reason"] = "collection_command_failed"
            else:
                partial = _partial_related_path(root, relative_path)
                if partial is not None:
                    attempt["final_status"] = "partial"
                    attempt["likely_reason"] = "partial_or_temp_file_present"
                    attempt["partial_path"] = str(partial)
                elif destination.is_file() and not _is_partial_name(destination.name):
                    attempt["final_status"] = "success"
                    attempt["bytes_copied"] = destination.stat().st_size
                    attempt["likely_reason"] = "artifact_collected"
                else:
                    attempt["final_status"] = "missing"
                    attempt["likely_reason"] = "remote_missing_or_not_produced"
        attempt["duration_seconds"] = max(0.0, now() - started)
        report["attempts"].append(attempt)
        _write_report(report_path, _refresh_summary(report, now=now()))
    report["batches"][-1]["finished_at"] = now()
    _write_report(report_path, _refresh_summary(report, now=now()))
    return report


def required_collection_failures(report: Mapping[str, Any], *, paths: Iterable[str] | None = None) -> list[dict[str, Any]]:
    selected = set(paths) if paths is not None else None
    failures: list[dict[str, Any]] = []
    for raw_attempt in report.get("attempts", []):
        if not isinstance(raw_attempt, dict):
            continue
        if selected is not None and raw_attempt.get("expected_path") not in selected:
            continue
        if raw_attempt.get("required") and raw_attempt.get("final_status") != "success":
            failures.append(dict(raw_attempt))
    return failures


def _load_report(path: Path) -> dict[str, Any]:
    if path.is_file():
        try:
            data = json.loads(path.read_text())
        except json.JSONDecodeError:
            data = None
        if isinstance(data, dict):
            data.setdefault("schema_version", 1)
            data.setdefault("attempts", [])
            data.setdefault("batches", [])
            data.setdefault("summary", {})
            return data
    return {"schema_version": 1, "attempts": [], "batches": [], "summary": {}}


def _refresh_summary(report: dict[str, Any], *, now: float) -> dict[str, Any]:
    counts: dict[str, int] = {}
    required_missing: list[str] = []
    for raw_attempt in report.get("attempts", []):
        if not isinstance(raw_attempt, dict):
            continue
        status = str(raw_attempt.get("final_status", "unknown"))
        counts[status] = counts.get(status, 0) + 1
        if raw_attempt.get("required") and status != "success":
            expected = raw_attempt.get("expected_path")
            if isinstance(expected, str):
                required_missing.append(expected)
    report["summary"] = {
        "updated_at": now,
        "attempt_count": sum(counts.values()),
        "status_counts": dict(sorted(counts.items())),
        "required_uncollected": required_missing,
        "ok": not required_missing,
    }
    return report


def _write_report(path: Path, report: Mapping[str, Any]) -> None:
    path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")


def _validate_relative_path(relative_path: str) -> None:
    path = Path(relative_path)
    if path.is_absolute() or ".." in path.parts:
        raise ValueError(f"Artifact path must be relative and stay under artifact root: {relative_path}")


def _partial_related_path(local_dir: Path, relative_path: str) -> Path | None:
    destination = local_dir / relative_path
    if destination.exists() and _is_partial_name(destination.name):
        return destination
    for suffix in _PARTIAL_SUFFIXES:
        candidate = destination.with_name(f"{destination.name}{suffix}")
        if candidate.exists():
            return candidate
    for temp_dir in _RSYNC_TEMP_DIRS:
        candidate = local_dir / temp_dir / relative_path
        if candidate.exists():
            return candidate
    return None


def _is_partial_name(name: str) -> bool:
    return name.endswith(_PARTIAL_SUFFIXES) or name in _RSYNC_TEMP_DIRS