"""Save decision-head weights and validate model and calibration identities."""

import hashlib
from pathlib import Path

import torch

from .provenance import RUNTIME_CONTRACT
from .atomic import atomic_output


def learned_state(model):
    return {"head": {key: tensor.detach().cpu().clone() for key, tensor in model.head.state_dict().items()}}


def learned_digest(model):
    digest = hashlib.sha256()
    for module, state in learned_state(model).items():
        for name, tensor in sorted(state.items()):
            digest.update(f"{module}.{name}:{tensor.dtype}:{tuple(tensor.shape)}".encode())
            digest.update(tensor.contiguous().numpy().tobytes())
    return digest.hexdigest()


def save_checkpoint(path, model, training=None):
    """Atomically save learned weights and optional resumable training state."""
    temperature = float(model.temperature)
    if not 0 < temperature < float("inf"):
        raise ValueError("Invalid checkpoint temperature")
    if model.calibration is None and temperature != 1.0:
        raise ValueError("A non-unit temperature needs outcome-fitted calibration provenance")
    if model.calibration and model.calibration["learned_sha256"] != learned_digest(model):
        raise ValueError("Calibration does not match current decision-head weights")
    if model.calibration and abs(model.calibration["temperature"]-temperature) > 1e-7:
        raise ValueError("Calibration temperature changed")
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {"contract": RUNTIME_CONTRACT, "identity": model.identity(),
               "learned": learned_state(model), "temperature": temperature,
               "calibration": model.calibration, "training": training}
    with atomic_output(path) as output:
        torch.save(payload, output)


def load_checkpoint(path, model):
    payload = torch.load(path, map_location="cpu", weights_only=True)
    if payload.get("contract") != RUNTIME_CONTRACT or payload.get("identity") != model.identity():
        raise ValueError("Checkpoint is not the exact current architecture/assets/preparation")
    if set(payload.get("learned", {})) != {"head"}:
        raise ValueError("Checkpoint must contain only decision-head weights")
    model.head.load_state_dict(payload["learned"]["head"], strict=True)
    temperature = float(payload["temperature"])
    if not 0 < temperature < float("inf"):
        raise ValueError("Invalid checkpoint temperature")
    calibration = payload.get("calibration")
    if calibration is None and temperature != 1.0:
        raise ValueError("A non-unit temperature needs outcome-fitted calibration provenance")
    if calibration and (calibration["learned_sha256"] != learned_digest(model)
                        or abs(calibration["temperature"]-temperature) > 1e-7):
        raise ValueError("Calibration binding differs from the saved weights/temperature")
    model.temperature.fill_(temperature)
    model.calibration = calibration
    model.eval()
    return payload