"""Validate the narrow Sudoku release surface before any training starts."""

from __future__ import annotations

import os
from pathlib import Path

import torch
from omegaconf import DictConfig
from softasp_core.training.schedule import CORRECTED_PROTOCOL
from torch import nn


def load_backbone(model: nn.Module, path: Path) -> None:
    """Load a complete extracted backbone, rejecting silent partial initialization."""
    if not path.is_file():
        raise FileNotFoundError(f"Existing backbone required: {path}; pretraining is not automatic")
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    state = checkpoint.get("model_state_dict") if isinstance(checkpoint, dict) else None
    expected = {k: v for k, v in model.state_dict().items() if k.startswith("cell_cnn.")}
    if not expected or not isinstance(state, dict) or set(state) != set(expected):
        raise ValueError("Backbone must contain exactly the complete cell_cnn.* tensor set")
    for key, target in expected.items():
        value = state[key]
        if (
            not isinstance(value, torch.Tensor)
            or value.shape != target.shape
            or value.dtype != target.dtype
            or not torch.isfinite(value).all()
        ):
            raise ValueError(f"Incompatible backbone tensor: {key}")
    model.load_state_dict(state, strict=False)


def prepare_release_model(cfg: DictConfig) -> nn.Module:
    """Reject unsupported overrides and load existing weights without training."""
    from softasp_solver.registry import get_model

    supported = (
        cfg.dataset.name == "sudoku"
        and cfg.trainer.get("protocol") == CORRECTED_PROTOCOL
        and cfg.model.name == "softasp_sudoku"
        and cfg.trainer.seed == 42
        and cfg.trainer.distributed == "false"
        and not cfg.trainer.compile
        and cfg.trainer.gradient_accumulation_steps == 1
        and cfg.trainer.resume_from is None
        and not cfg.model.params.get("oracle_perception", False)
        and int(os.environ.get("WORLD_SIZE", "1")) == 1
    )
    if not supported:
        raise ValueError(
            "Release 1 requires Sudoku AS2, seed 42, single device, no oracle/resume/compile"
        )
    backbone = cfg.model.get("pretrained_from")
    if not backbone:
        raise ValueError(
            "Set model.pretrained_from to an existing backbone; no automatic pretraining"
        )
    model = get_model(cfg.model.name, **dict(cfg.model.params))
    load_backbone(model, Path(backbone))
    return model