"""Validated backbone handoff between the two Sudoku training stages."""

from pathlib import Path

import torch

from softasp.sudoku_reporting import file_sha256


def extract_backbone(checkpoint: Path, destination: Path) -> dict:
    """Extract a newly trained perception checkpoint, retaining source identity."""
    saved = torch.load(checkpoint, map_location="cpu", weights_only=False)
    if saved.get("config", {}).get("model", {}).get("name") != "perception_only_sudoku":
        raise ValueError("Expected a perception_only_sudoku training checkpoint")
    if saved.get("epoch", -1) < 0:
        raise ValueError("Backbone checkpoint must follow at least one training epoch")
    state = {
        key: value
        for key, value in saved["model_state_dict"].items()
        if key.startswith("cell_cnn.")
    }
    if not state or any(not torch.isfinite(value).all() for value in state.values()):
        raise ValueError("Expected finite trained cell_cnn tensors")
    provenance = {
        "source_checkpoint": str(checkpoint.resolve()),
        "source_sha256": file_sha256(checkpoint),
        "source_epoch": saved["epoch"],
        "source_metric": saved.get("best_metric"),
        "dataset_identity": saved["config"].get("resolved_dataset"),
    }
    # Refuse to overwrite any pre-existing model artifact.
    with destination.open("xb") as stream:
        torch.save({"model_state_dict": state, **provenance}, stream)
    return {**provenance, "path": str(destination), "sha256": file_sha256(destination)}