"""Explicit local inference selection without importing optional ML dependencies."""

import importlib
from dataclasses import dataclass
from enum import StrEnum
from pathlib import Path
from typing import Protocol, cast

from .suggestion import SuggestionResult, suggest

_OPTIONAL_MODULES = frozenset({"numpy", "packaging", "safetensors", "torch"})


class InferenceError(ValueError):
    """An explicit inference backend cannot be configured or loaded safely."""


class InferenceBackend(StrEnum):
    """Available instruction interpreters; baseline remains the default."""

    BASELINE = "baseline"
    LEARNED = "learned"


def _validate_device(device: object) -> None:
    """Reject invalid or unhashable device values at the runtime boundary."""
    if not isinstance(device, str) or device not in {"cpu", "mps"}:
        raise InferenceError("Device must be explicitly 'cpu' or 'mps'.")


@dataclass(frozen=True, slots=True)
class InferenceConfig:
    """Explicit backend, acknowledgement, checkpoint, and device selection."""

    backend: InferenceBackend = InferenceBackend.BASELINE
    allow_experimental: bool = False
    config_path: Path | None = None
    weights_path: Path | None = None
    device: str = "cpu"

    def __post_init__(self) -> None:
        """Reject ignored or incomplete backend options."""
        if type(self.backend) is not InferenceBackend:
            raise InferenceError("Backend must be an InferenceBackend enum value.")
        if type(self.allow_experimental) is not bool:
            raise InferenceError("Experimental acknowledgement must be an exact boolean.")
        if self.config_path is not None and not isinstance(self.config_path, Path):
            raise InferenceError("Config checkpoint must be a Path or None.")
        if self.weights_path is not None and not isinstance(self.weights_path, Path):
            raise InferenceError("Weights checkpoint must be a Path or None.")
        _validate_device(self.device)
        learned_options = (
            self.allow_experimental or self.config_path is not None or self.weights_path is not None
        )
        if self.backend == InferenceBackend.BASELINE:
            if learned_options or self.device != "cpu":
                raise InferenceError(
                    "Checkpoint, device, and opt-in options require learned backend."
                )
            return
        if not self.allow_experimental:
            raise InferenceError("Learned inference requires allow_experimental=True.")
        if self.config_path is None or self.weights_path is None:
            raise InferenceError("Learned inference requires explicit config and weights paths.")


class LearnedRuntime(Protocol):
    """Minimal typed boundary implemented by the optional learned runtime."""

    @property
    def device(self) -> str:
        """Return the device actually selected without fallback."""
        ...

    def __call__(self, source: str, instruction: str) -> SuggestionResult:
        """Support the existing evaluation predictor contract."""
        ...

    def suggest(self, source: str, instruction: str) -> SuggestionResult:
        """Return one immutable core suggestion result."""
        ...


class _RuntimeType(Protocol):
    def load(self, config_path: Path, weights_path: Path, device: str) -> LearnedRuntime:
        """Load one runtime through the hardened optional implementation."""
        ...


def _load_learned_runtime(config_path: Path, weights_path: Path, device: str) -> LearnedRuntime:
    module = importlib.import_module(".ml.runtime", __package__)
    runtime_type = cast("_RuntimeType", module.__dict__["LoadedLearnedRuntime"])
    return runtime_type.load(config_path, weights_path, device)


def _optional_dependency_error(error: ModuleNotFoundError) -> InferenceError | None:
    name = error.name or ""
    top_level = name.partition(".")[0]
    if top_level not in _OPTIONAL_MODULES:
        return None
    return InferenceError(
        "Learned inference dependencies are unavailable; install the local package with "
        "the 'yaml-edit[ml]' extra."
    )


class InferenceEngine:
    """Reusable explicit inference adapter that loads at most one learned model."""

    def __init__(self, config: InferenceConfig | None = None) -> None:
        """Validate configuration and eagerly load one requested local checkpoint."""
        self.config = config if config is not None else InferenceConfig()
        self._runtime: LearnedRuntime | None = None
        if self.config.backend == InferenceBackend.LEARNED:
            self._runtime = self._load_runtime()

    def _load_runtime(self) -> LearnedRuntime:
        config_path = self.config.config_path
        weights_path = self.config.weights_path
        if config_path is None or weights_path is None:
            raise InferenceError("Learned inference requires explicit checkpoint paths.")
        try:
            return _load_learned_runtime(config_path, weights_path, self.config.device)
        except ModuleNotFoundError as error:
            classified = _optional_dependency_error(error)
            if classified is None:
                raise
            raise classified from error
        except (OSError, ValueError) as error:
            raise InferenceError(
                f"Learned inference could not load the checkpoint: {error}"
            ) from error

    @property
    def runtime(self) -> LearnedRuntime | None:
        """Expose the one explicitly loaded runtime for evaluation and measurement."""
        return self._runtime

    def suggest(self, source: str, instruction: str) -> SuggestionResult:
        """Interpret without applying or writing the source document."""
        if self._runtime is None:
            return suggest(source, instruction)
        return self._runtime.suggest(source, instruction)