src/provenance.py

"""
Provenance, content-hashing, and dataset-reuse prevention.

Design goals (per project rules):
  - Single source of truth for "has this dataset already been trained on".
  - No silent fallback behavior anywhere in this module.
  - Dataset claims are atomic (safe under concurrent/multi-contributor use).
  - Checkpoints are content-addressed (Option A): filename = hash(config+code+dataset).

This module intentionally has NO knowledge of model internals. It only
handles identity, claiming, and structured errors.
"""
from __future__ import annotations
import hashlib
import json
import os
import time
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any, Dict, Optional

import torch


# ----------------------------- Errors ----------------------------------- #

class ProvenanceError(Exception):
    outcome_code: str = "PROVENANCE_ERROR"

    def __init__(self, detail: str, outcome_code: Optional[str] = None):
        self.detail = detail
        if outcome_code:
            self.outcome_code = outcome_code
        super().__init__(f"[{self.outcome_code}] {detail}")


class DataLoadError(ProvenanceError):
    outcome_code = "DATA_LOAD_FAILED"


class SchemaValidationError(ProvenanceError):
    outcome_code = "SCHEMA_VALIDATION_FAILED"


class DatasetAlreadyUsedError(ProvenanceError):
    outcome_code = "ALREADY_CONSUMED"


class DatasetInProgressError(ProvenanceError):
    outcome_code = "ALREADY_IN_PROGRESS"


class CheckpointIntegrityError(ProvenanceError):
    outcome_code = "CHECKPOINT_INTEGRITY_FAILED"


class TrajectoryTooShortError(ProvenanceError):
    outcome_code = "TRAJECTORY_TOO_SHORT"


class EmptyDatasetError(ProvenanceError):
    outcome_code = "EMPTY_DATASET"


def validate_trajectory_lengths(
    dataset,
    required_length: int,
    sample_cap: int = 64,
):
    n = len(dataset)
    if n == 0:
        raise EmptyDatasetError(
            "dataset has zero trajectories; nothing to train on",
        )
    check_n = min(n, sample_cap)
    too_short = []
    for i in range(check_n):
        item = dataset[i]
        fields = item["fields"] if isinstance(item, dict) else item
        T = fields.shape[0]
        if T < required_length:
            too_short.append((i, T))
    if too_short:
        raise TrajectoryTooShortError(
            f"{len(too_short)}/{check_n} sampled trajectories are shorter "
            f"than required_length={required_length} "
            f"(examples: {too_short[:5]}). Every trajectory must satisfy "
            f"T >= window + pred_steps; silently skipping short "
            f"trajectories mid-training is not permitted.",
        )
    return {"success": True, "outcome_code": "OK", "checked": check_n, "total": n}


# ----------------------------- Hashing ----------------------------------- #

def hash_config(config: Dict[str, Any]) -> str:
    def _norm(v):
        if isinstance(v, float):
            return round(v, 10)
        if isinstance(v, dict):
            return {k: _norm(vv) for k, vv in sorted(v.items())}
        if isinstance(v, (list, tuple)):
            return [_norm(vv) for vv in v]
        return v

    normalized = _norm(config)
    blob = json.dumps(normalized, sort_keys=True).encode("utf-8")
    return hashlib.sha256(blob).hexdigest()


def hash_code(src_dir: str) -> str:
    src_path = Path(src_dir)
    py_files = sorted(src_path.rglob("*.py"))
    hasher = hashlib.sha256()
    for f in py_files:
        hasher.update(f.name.encode("utf-8"))
        hasher.update(f.read_bytes())
    return hasher.hexdigest()


def hash_dataset(dataset, sample_cap: Optional[int] = None) -> str:
    n = len(dataset)
    if sample_cap is not None:
        n = min(n, sample_cap)
    hasher = hashlib.sha256()
    hasher.update(str(n).encode("utf-8"))
    for i in range(n):
        item = dataset[i]
        fields = item["fields"] if isinstance(item, dict) else item
        if not torch.is_tensor(fields):
            raise SchemaValidationError(
                f"dataset[{i}] did not return a tensor under key 'fields'; "
                f"got {type(fields)}"
            )
        arr = fields.detach().cpu().contiguous().numpy()
        hasher.update(arr.tobytes())
        hasher.update(str(arr.shape).encode("utf-8"))
    return hasher.hexdigest()


def combined_identity_hash(config_hash: str, code_hash: str, dataset_hash: str) -> str:
    blob = f"{config_hash}:{code_hash}:{dataset_hash}".encode("utf-8")
    return hashlib.sha256(blob).hexdigest()


# ------------------------- Dataset reuse registry ------------------------ #

@dataclass
class ClaimResult:
    success: bool
    outcome_code: str
    dataset_hash: str
    detail: str = ""


class DatasetRegistry:
    def __init__(self, registry_dir: str = "registry/datasets"):
        self.dir = Path(registry_dir)
        self.dir.mkdir(parents=True, exist_ok=True)

    def _path(self, dataset_hash: str) -> Path:
        return self.dir / f"{dataset_hash}.json"

    def status(self, dataset_hash: str) -> Optional[Dict[str, Any]]:
        p = self._path(dataset_hash)
        if not p.exists():
            return None
        return json.loads(p.read_text())

    def claim(self, dataset_hash: str, experiment_id: str) -> ClaimResult:
        existing = self.status(dataset_hash)
        if existing is not None:
            if existing["status"] == "CONSUMED":
                raise DatasetAlreadyUsedError(
                    f"dataset {dataset_hash[:12]} was already consumed by "
                    f"experiment {existing.get('experiment_id')} at "
                    f"{existing.get('consumed_at')}. Retraining on it is blocked."
                )
            if existing["status"] == "IN_PROGRESS":
                raise DatasetInProgressError(
                    f"dataset {dataset_hash[:12]} is currently IN_PROGRESS "
                    f"(experiment {existing.get('experiment_id')}, claimed "
                    f"{existing.get('claimed_at')}). If that run crashed, "
                    f"resolve manually with DatasetRegistry.mark_failed() "
                    f"before retrying — this is not automatic by design."
                )
            if existing["status"] == "FAILED":
                raise DatasetInProgressError(
                    f"dataset {dataset_hash[:12]} previously FAILED "
                    f"(experiment {existing.get('experiment_id')}). "
                    f"Explicit human confirmation required to retry: "
                    f"call DatasetRegistry.allow_retry(dataset_hash) first."
                )

        record = {
            "status": "IN_PROGRESS",
            "experiment_id": experiment_id,
            "claimed_at": time.time(),
            "consumed_at": None,
        }
        p = self._path(dataset_hash)
        try:
            # O_EXCL makes this atomic: fails if another process just created it.
            fd = os.open(str(p), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
            with os.fdopen(fd, "w") as f:
                json.dump(record, f, indent=2)
        except FileExistsError:
            # lost the race — re-check what the winner wrote
            return self.claim(dataset_hash, experiment_id)
        return ClaimResult(True, "CLAIMED", dataset_hash)

    def mark_consumed(self, dataset_hash: str):
        record = self.status(dataset_hash)
        if record is None:
            raise ProvenanceError(
                f"cannot mark {dataset_hash[:12]} consumed: no claim exists",
                outcome_code="NO_CLAIM_FOUND",
            )
        record["status"] = "CONSUMED"
        record["consumed_at"] = time.time()
        self._path(dataset_hash).write_text(json.dumps(record, indent=2))

    def mark_failed(self, dataset_hash: str, error_detail: str = ""):
        record = self.status(dataset_hash)
        if record is None:
            return
        record["status"] = "FAILED"
        record["error_detail"] = error_detail
        self._path(dataset_hash).write_text(json.dumps(record, indent=2))

    def allow_retry(self, dataset_hash: str):
        """Explicit human action required to clear a FAILED claim. Not automatic."""
        p = self._path(dataset_hash)
        if p.exists():
            p.unlink()


# ------------------------------ Checkpoints ------------------------------- #

def _json_safe(obj: Any) -> Any:
    if isinstance(obj, torch.Tensor):
        return obj.detach().cpu().tolist()
    if isinstance(obj, dict):
        return {k: _json_safe(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_json_safe(v) for v in obj]
    if isinstance(obj, (str, int, float, bool)) or obj is None:
        return obj
    try:
        json.dumps(obj)
        return obj
    except TypeError:
        return str(obj)


class CheckpointStore:
    def __init__(self, checkpoints_dir: str = "checkpoints"):
        self.dir = Path(checkpoints_dir)
        self.dir.mkdir(parents=True, exist_ok=True)

    def save(
        self,
        model_state: Dict[str, Any],
        config: Dict[str, Any],
        dataset_hash: str,
        code_hash: str,
        data_provenance: str,
        extra: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        if not dataset_hash:
            raise ProvenanceError(
                "refusing to save checkpoint without dataset_hash",
                outcome_code="MISSING_DATASET_HASH",
            )
        if data_provenance not in ("REAL_LOCAL", "REAL_STREAMED", "REAL_PBDB",
                                     "REAL_PBDB_TAXONOMY", "SYNTHETIC_TREE", "SYNTHETIC"):
            raise ProvenanceError(
                f"invalid data_provenance '{data_provenance}'",
                outcome_code="INVALID_PROVENANCE",
            )

        config_hash = hash_config(config)
        identity = combined_identity_hash(config_hash, code_hash, dataset_hash)
        final_path = self.dir / f"{identity}.pt"
        meta_path = self.dir / f"{identity}.meta.json"

        if final_path.exists() and meta_path.exists():
            return {
                "success": True,
                "outcome_code": "DUPLICATE_EXISTS",
                "path": str(final_path),
                "identity_hash": identity,
            }

        meta = {
            "identity_hash": identity,
            "config_hash": config_hash,
            "code_hash": code_hash,
            "dataset_hash": dataset_hash,
            "data_provenance": data_provenance,
            "config": _json_safe(config),
            "created_at": time.time(),
            "extra": _json_safe(extra or {}),
        }
        meta_json_str = json.dumps(meta, indent=2)  # _json_safe guarantees this succeeds

        pid_tag = f"{os.getpid()}_{int(time.time()*1000)}"
        tmp_path = self.dir / f".tmp_{identity}_{pid_tag}.pt"
        meta_tmp = self.dir / f".tmp_{identity}_meta_{pid_tag}.json"
        try:
            torch.save(model_state, tmp_path)
            meta_tmp.write_text(meta_json_str)
            os.replace(tmp_path, final_path)
            os.replace(meta_tmp, meta_path)
        finally:
            # Clean up any tmp file left behind by a failed/partial attempt.
            for p in (tmp_path, meta_tmp):
                if p.exists():
                    p.unlink()

        return {
            "success": True,
            "outcome_code": "SAVED",
            "path": str(final_path),
            "identity_hash": identity,
        }

    def load(self, identity_hash: str) -> Dict[str, Any]:
        final_path = self.dir / f"{identity_hash}.pt"
        meta_path = self.dir / f"{identity_hash}.meta.json"
        if not final_path.exists() or not meta_path.exists():
            raise CheckpointIntegrityError(
                f"checkpoint {identity_hash[:12]} incomplete or missing "
                f"(model or meta file absent)"
            )
        model_state = torch.load(final_path, map_location="cpu")
        meta = json.loads(meta_path.read_text())
        return {"model_state": model_state, "meta": meta}