"""Default-off learned predictor that can only compile core-valid copied actions."""

from dataclasses import dataclass

import torch

from yaml_edit.engine import preview, propose

from .checkpoint import ArtifactConfig
from .features import MlInputError, prepare_input
from .model import CandidateScorer, collate_features

_DISPOSITIONS = ("supported", "needs_clarification", "unsupported")


@dataclass(slots=True)
class LearnedPredictor:
    """Callable evaluator-compatible wrapper around a loaded candidate scorer."""

    config: ArtifactConfig
    model: CandidateScorer
    device: str = "cpu"

    def __post_init__(self) -> None:
        """Require a coherent checkpoint and explicit device without fallback."""
        self.config.validate()
        self.model.validate_architecture(
            self.config.feature_dimension, self.config.hidden_dimension
        )
        if self.device not in {"cpu", "mps"}:
            raise ValueError("Device must be explicitly 'cpu' or 'mps'.")
        if self.device == "mps" and not torch.backends.mps.is_available():
            raise ValueError("MPS was requested but is not available.")
        self.model.to(torch.device(self.device))
        self.model.eval()

    def __call__(self, source: str, instruction: str) -> dict[str, object]:
        """Predict from source/instruction only, abstaining if no valid action exists."""
        try:
            prepared = prepare_input(source, instruction)
        except MlInputError:
            return {"disposition": "unsupported", "proposal": None}
        batch = collate_features(
            [list(prepared.vector)],
            [[list(action.features) for action in prepared.actions]],
        )
        device = torch.device(self.device)
        with torch.no_grad():
            disposition_logits, action_logits = self.model(
                batch.vectors.to(device),
                batch.actions.to(device),
                batch.action_mask.to(device),
            )
            disposition_probabilities = torch.softmax(disposition_logits[0], dim=0)
        disposition_index = int(torch.argmax(disposition_probabilities).item())
        supported_confidence = float(disposition_probabilities[0].item())
        if disposition_index == 0 and supported_confidence < self.config.disposition_threshold:
            disposition_index = int(torch.argmax(disposition_probabilities[1:]).item()) + 1
        if disposition_index != 0:
            return {"disposition": _DISPOSITIONS[disposition_index], "proposal": None}
        if not prepared.actions:
            return {"disposition": "unsupported", "proposal": None}
        with torch.no_grad():
            action_probabilities = torch.softmax(action_logits[0, : len(prepared.actions)], dim=0)
        action_index = int(torch.argmax(action_probabilities).item())
        if float(action_probabilities[action_index].item()) < self.config.action_threshold:
            return {"disposition": "unsupported", "proposal": None}
        action = prepared.actions[action_index]
        proposal = propose(source, action.candidate.id, action.value)
        preview(source, proposal)
        return {"disposition": "supported", "proposal": proposal.to_dict()}