"""Reloadable inference artifact export."""

from __future__ import annotations

import json
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import torch

from graph.config import GraphConfig, PATConfig
from graph.model import GraphModel


def load_checkpoint_file(
    path: str | Path, *, map_location: str | torch.device = "cpu"
) -> dict[str, Any]:
    """Load a GRM checkpoint with ``weights_only=True``.

    Checkpoints written by :class:`~graph.training.trainer.GraphTrainer` contain
    only tensors and plain containers, so unpickling arbitrary Python objects is
    never required.  Loading model artifacts with ``weights_only=False`` would
    permit code execution from a checkpoint file, which matters because the
    registry promotion flow (PLAN section 40) moves artifacts between machines.
    """
    # A checkpoint is a heterogeneous payload from outside the process; this is
    # the external boundary where dict[str, Any] is the honest type.
    payload = torch.load(path, map_location=map_location, weights_only=True)
    if not isinstance(payload, dict):
        raise TypeError(f"{path}: checkpoint is not a dict payload")
    return payload


def save_inference_artifact(
    model: GraphModel, directory: str | Path, *, use_safetensors: bool = False
) -> Path:
    root = Path(directory)
    root.mkdir(parents=True, exist_ok=True)
    if use_safetensors:
        try:
            from safetensors.torch import save_file
        except (
            ImportError
        ) as exc:  # pragma: no cover - depends on optional export dependency
            raise RuntimeError(
                "safetensors export requested but safetensors is not installed"
            ) from exc
        weights = root / "model.safetensors"
        save_file(
            {key: value.detach().cpu() for key, value in model.state_dict().items()},
            str(weights),
        )
        weight_format = "safetensors"
    else:
        weights = root / "model.pt"
        torch.save(model.state_dict(), weights)
        weight_format = "torch"
    config = model.config.to_dict()
    config["weight_format"] = weight_format
    (root / "config.json").write_text(
        json.dumps(config, indent=2, sort_keys=True), encoding="utf-8"
    )
    return weights


def load_inference_artifact(
    directory: str | Path, *, map_location: str | torch.device = "cpu"
) -> GraphModel:
    root = Path(directory)
    raw = json.loads((root / "config.json").read_text(encoding="utf-8"))
    weight_format = raw.pop("weight_format", "torch")
    pat = PATConfig(**raw.pop("pat"))
    raw.pop("training", None)
    raw.pop("loss", None)
    config = GraphConfig(pat=pat, **raw)
    model = GraphModel(config)
    if weight_format == "safetensors":
        try:
            from safetensors.torch import load_file
        except (
            ImportError
        ) as exc:  # pragma: no cover - depends on optional export dependency
            raise RuntimeError(
                "safetensors artifact requires the safetensors package"
            ) from exc
        state = load_file(str(root / "model.safetensors"), device=str(map_location))
    else:
        state = torch.load(
            root / "model.pt", map_location=map_location, weights_only=True
        )
    load_state_dict_compat(model, state)
    model.eval()
    return model


def load_state_dict_compat(
    model: GraphModel, state: Mapping[str, torch.Tensor]
) -> None:
    """Load checkpoints across explicitly supported PAT shape migrations.

    The structured PAT token expansion added type/attribute tables and a
    larger position table.  Older registered smoke artifacts remain
    loadable by zero-initializing only those new inputs and padding the old
    position rows; every other missing or incompatible tensor fails loudly.
    """
    expected = model.state_dict()
    migrated: dict[str, torch.Tensor] = {}
    missing_allowed = {
        "pat.type_embedding.weight",
        "pat.attribute_embedding.weight",
    }
    for key, value in state.items():
        if key not in expected:
            raise RuntimeError(f"checkpoint contains unexpected tensor: {key}")
        if value.shape == expected[key].shape:
            migrated[key] = value
            continue
        if (
            key == "pat.path_position_embedding.weight"
            and value.ndim == 2
            and value.shape[1] == expected[key].shape[1]
            and value.shape[0] < expected[key].shape[0]
        ):
            padded = torch.zeros_like(expected[key])
            padded[: value.shape[0]] = value
            migrated[key] = padded
            continue
        raise RuntimeError(
            f"checkpoint tensor shape mismatch for {key}: {tuple(value.shape)} != {tuple(expected[key].shape)}"
        )
    for key, value in expected.items():
        if key not in migrated:
            if key in missing_allowed:
                migrated[key] = torch.zeros_like(value)
            else:
                raise RuntimeError(f"checkpoint is missing tensor: {key}")
    model.load_state_dict(migrated, strict=True)