"""Versioned JSON plus non-executable safetensors checkpoint contract."""

import json
import math
import os
import tempfile
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Never, cast

import torch
from safetensors import SafetensorError, safe_open
from safetensors.torch import save_file

from .features import ACTION_FEATURES, FEATURE_DIMENSION
from .model import DISPOSITION_COUNT, MAX_HIDDEN_DIMENSION, CandidateScorer

CHECKPOINT_FORMAT = "yaml-edit/ml-checkpoint-v1"
MAX_CONFIG_BYTES = 64 * 1024
MAX_WEIGHTS_BYTES = 10 * 1024 * 1024
_SHA256_LENGTH = 64
_MAX_SEED = 2**31 - 1
_MAX_EXAMPLES = 10_000
_MAX_EPOCH = 100_000
_CONFIG_KEYS = frozenset(
    {
        "format",
        "feature_dimension",
        "action_features",
        "hidden_dimension",
        "disposition_threshold",
        "action_threshold",
        "seed",
        "dataset_sha256",
        "train_examples",
        "validation_examples",
        "selected_epoch",
    }
)


class CheckpointError(ValueError):
    """A checkpoint violates the bounded safe-artifact contract."""


def _fail(message: str) -> Never:
    raise CheckpointError(message)


@dataclass(frozen=True, slots=True)
class ArtifactConfig:
    """All architecture and calibrated inference settings needed for loading."""

    feature_dimension: int
    hidden_dimension: int
    disposition_threshold: float
    action_threshold: float
    seed: int
    dataset_sha256: str
    train_examples: int
    validation_examples: int
    selected_epoch: int

    def __post_init__(self) -> None:
        """Reject incoherent in-memory configurations at their construction boundary."""
        self.validate()

    def validate(self) -> None:
        """Validate exact types, fixed architecture, bounds, and finite thresholds."""
        _bounded_int(
            self.feature_dimension, "feature_dimension", FEATURE_DIMENSION, FEATURE_DIMENSION
        )
        _bounded_int(self.hidden_dimension, "hidden_dimension", 1, MAX_HIDDEN_DIMENSION)
        _threshold(self.disposition_threshold, "disposition_threshold")
        _threshold(self.action_threshold, "action_threshold")
        _bounded_int(self.seed, "seed", 0, _MAX_SEED)
        _bounded_int(self.train_examples, "train_examples", 1, _MAX_EXAMPLES)
        _bounded_int(self.validation_examples, "validation_examples", 1, _MAX_EXAMPLES)
        _bounded_int(self.selected_epoch, "selected_epoch", 0, _MAX_EPOCH)
        _digest(self.dataset_sha256)

    def to_dict(self) -> dict[str, object]:
        """Return the strict versioned JSON record."""
        self.validate()
        result: dict[str, object] = dict(asdict(self))
        result.update({"format": CHECKPOINT_FORMAT, "action_features": ACTION_FEATURES})
        return result

    @classmethod
    def from_dict(cls, value: object) -> "ArtifactConfig":
        """Validate untrusted decoded configuration before model allocation."""
        if not isinstance(value, dict) or set(value) != _CONFIG_KEYS:
            _fail("Checkpoint config has missing or unknown fields.")
        data = cast("dict[str, object]", value)
        if data["format"] != CHECKPOINT_FORMAT:
            _fail("Checkpoint format is unsupported.")
        _bounded_int(data["action_features"], "action width", ACTION_FEATURES, ACTION_FEATURES)
        return cls(
            _bounded_int(
                data["feature_dimension"],
                "feature_dimension",
                FEATURE_DIMENSION,
                FEATURE_DIMENSION,
            ),
            _bounded_int(data["hidden_dimension"], "hidden_dimension", 1, MAX_HIDDEN_DIMENSION),
            _threshold(data["disposition_threshold"], "disposition_threshold"),
            _threshold(data["action_threshold"], "action_threshold"),
            _bounded_int(data["seed"], "seed", 0, _MAX_SEED),
            _digest(data["dataset_sha256"]),
            _bounded_int(data["train_examples"], "train_examples", 1, _MAX_EXAMPLES),
            _bounded_int(data["validation_examples"], "validation_examples", 1, _MAX_EXAMPLES),
            _bounded_int(data["selected_epoch"], "selected_epoch", 0, _MAX_EPOCH),
        )


def _bounded_int(value: object, label: str, minimum: int, maximum: int) -> int:
    if type(value) is not int or not minimum <= value <= maximum:
        _fail(f"{label} is outside its integer bounds.")
    return value


def _threshold(value: object, label: str) -> float:
    if type(value) not in (int, float):
        _fail(f"{label} must be finite and within [0, 1].")
    numeric = cast("int | float", value)
    if not math.isfinite(numeric) or not 0.0 <= numeric <= 1.0:
        _fail(f"{label} must be finite and within [0, 1].")
    return float(numeric)


def _digest(value: object) -> str:
    if not isinstance(value, str) or len(value) != _SHA256_LENGTH:
        _fail("dataset_sha256 must be a 64-character digest.")
    if any(character not in "0123456789abcdef" for character in value):
        _fail("dataset_sha256 must be lowercase hexadecimal.")
    return value


def _expected_shapes(config: ArtifactConfig) -> dict[str, tuple[int, ...]]:
    hidden = config.hidden_dimension
    return {
        "encoder.weight": (hidden, config.feature_dimension),
        "encoder.bias": (hidden,),
        "disposition.weight": (DISPOSITION_COUNT, hidden),
        "disposition.bias": (DISPOSITION_COUNT,),
        "action.weight": (hidden, ACTION_FEATURES),
        "action.bias": (hidden,),
    }


def _read_capped(path: Path, limit: int, label: str) -> bytes:
    try:
        with path.open("rb") as file:
            raw = file.read(limit + 1)
    except OSError as error:
        raise CheckpointError(f"Checkpoint {label} cannot be read.") from error
    if len(raw) > limit:
        _fail(f"Checkpoint {label} size exceeds the hard limit.")
    return raw


def _pairs(items: list[tuple[str, object]]) -> dict[str, object]:
    result: dict[str, object] = {}
    for key, value in items:
        if key in result:
            _fail(f"Duplicate JSON key: {key!r}.")
        result[key] = value
    return result


def _load_json(path: Path) -> ArtifactConfig:
    raw = _read_capped(path, MAX_CONFIG_BYTES, "config")
    try:
        value: object = json.loads(
            raw.decode("utf-8"), object_pairs_hook=_pairs, parse_constant=_nonfinite
        )
    except (UnicodeError, json.JSONDecodeError, RecursionError) as error:
        raise CheckpointError("Checkpoint config is not valid bounded JSON.") from error
    return ArtifactConfig.from_dict(value)


def _nonfinite(value: str) -> Never:
    _fail(f"Nonfinite JSON number is forbidden: {value}.")


def _validate_model(config: ArtifactConfig, model: CandidateScorer) -> None:
    config.validate()
    expected = _expected_shapes(config)
    state = model.state_dict()
    if set(state) != set(expected):
        _fail("Checkpoint model keys do not match the model contract.")
    for key, shape in expected.items():
        tensor = state[key]
        if tuple(tensor.shape) != shape or tensor.dtype != torch.float32:
            _fail("Checkpoint model is incompatible with its configuration.")
        if not bool(torch.isfinite(tensor).all()):
            _fail("Checkpoint tensors must contain only finite values.")


def _load_snapshot(raw: bytes, config: ArtifactConfig) -> CandidateScorer:
    expected = _expected_shapes(config)
    try:
        with tempfile.TemporaryDirectory(prefix="yaml-edit-weights-") as directory:
            snapshot = Path(directory) / "weights.safetensors"
            snapshot.write_bytes(raw)
            snapshot.chmod(0o600)
            with safe_open(  # type: ignore[no-untyped-call]
                snapshot, framework="pt", device="cpu"
            ) as tensors:
                if set(tensors.keys()) != set(expected):
                    _fail("Checkpoint tensor keys do not match the model contract.")
                for key, shape in expected.items():
                    tensor_slice = tensors.get_slice(key)
                    if tensor_slice.get_shape() != list(shape):
                        _fail(f"Checkpoint tensor {key!r} has the wrong shape.")
                    if tensor_slice.get_dtype() != "F32":
                        _fail(f"Checkpoint tensor {key!r} has the wrong dtype.")
                model = CandidateScorer(config.feature_dimension, config.hidden_dimension)
                state = {key: tensors.get_tensor(key) for key in expected}
    except CheckpointError:
        raise
    except OSError as error:
        raise CheckpointError("Checkpoint weight snapshot cannot be created.") from error
    except SafetensorError as error:
        raise CheckpointError("Checkpoint safetensors header or data is invalid.") from error
    if any(not bool(torch.isfinite(tensor).all()) for tensor in state.values()):
        _fail("Checkpoint tensors must contain only finite values.")
    model.load_state_dict(state, strict=True)
    model.eval()
    return model


def load_checkpoint(
    config_path: Path, weights_path: Path
) -> tuple[ArtifactConfig, CandidateScorer]:
    """Load one bounded immutable snapshot after strict pre-allocation validation."""
    config = _load_json(config_path)
    weights = _read_capped(weights_path, MAX_WEIGHTS_BYTES, "weights")
    return config, _load_snapshot(weights, config)


def save_checkpoint(
    config_path: Path,
    weights_path: Path,
    config: ArtifactConfig,
    model: CandidateScorer,
) -> None:
    """Privately publish a validated checkpoint pair without overwriting either path."""
    _validate_model(config, model)
    if config_path.exists() or weights_path.exists():
        raise FileExistsError("Checkpoint output paths must not already exist.")
    config_bytes = (
        json.dumps(config.to_dict(), indent=2, sort_keys=True, allow_nan=False) + "\n"
    ).encode("utf-8")
    if len(config_bytes) > MAX_CONFIG_BYTES:
        _fail("Checkpoint config size exceeds the hard limit.")
    state = {key: value.detach().cpu().contiguous() for key, value in model.state_dict().items()}
    with tempfile.TemporaryDirectory(dir=config_path.parent) as directory:
        temporary = Path(directory)
        temporary_weights = temporary / "weights.safetensors"
        temporary_config = temporary / "config.json"
        save_file(state, temporary_weights)
        if temporary_weights.stat().st_size > MAX_WEIGHTS_BYTES:
            _fail("Checkpoint weights size exceeds the hard limit.")
        temporary_config.write_bytes(config_bytes)
        temporary_weights.chmod(0o600)
        temporary_config.chmod(0o600)
        os.link(temporary_weights, weights_path)
        try:
            os.link(temporary_config, config_path)
        except OSError:
            weights_path.unlink(missing_ok=True)
            raise