"""Checkpoint provenance validation shared by instantiation adapters."""

from pathlib import Path

import numpy as np
import torch
from numpy._core.multiarray import _reconstruct
from omegaconf import DictConfig, OmegaConf

from softasp.instantiations.runtime import identity


def read_checkpoint(path: Path) -> dict:
    # Full trainer checkpoints contain NumPy RNG state. Never retry with
    # unrestricted pickle loading; allow only its concrete array representation.
    with torch.serialization.safe_globals(
        [_reconstruct, np.ndarray, np.dtype, type(np.dtype("uint32"))]
    ):
        return torch.load(path, map_location="cpu", weights_only=True)


def validate_checkpoint(
    cfg: DictConfig,
    saved: dict,
    *,
    legacy_model: str | None = None,
    legacy_constraints: dict | None = None,
) -> None:
    expected = identity(cfg)
    recorded = saved.get("config", {})
    actual = recorded.get("instantiation")
    selected = cfg.instantiation.get("constraints")
    constraints = OmegaConf.to_container(selected, resolve=True) if selected is not None else None
    if actual is None:
        if legacy_model is None or recorded.get("model", {}).get("name") != legacy_model:
            raise ValueError("Checkpoint has no compatible instantiation identity")
        if constraints != legacy_constraints:
            raise ValueError("Legacy checkpoint constraint provider/parameters do not match")
        return
    if any(actual.get(key) != value for key, value in expected.items()):
        raise ValueError("Checkpoint instantiation identity does not match configuration")
    if recorded.get("model", {}).get("name") != cfg.get("model", {}).get("name"):
        raise ValueError("Checkpoint model identity does not match configuration")
    if actual.get("constraints") != constraints:
        raise ValueError("Checkpoint constraint provider/parameters do not match configuration")
    network = cfg.instantiation.get("network")
    network = OmegaConf.to_container(network, resolve=True) if network is not None else None
    if actual.get("network") != network:
        raise ValueError("Checkpoint network configuration does not match")