"""Model loader helpers for eval runtime selection."""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass, field


ModelLoaderCallable = Callable[[object], object]


def load_model_with_runtime(
    action_tracker: object,
    runtime_framework: str,
    loaders: dict[str, ModelLoaderCallable],
    default_key: str = "pytorch",
) -> object:
    """Load evaluation model by runtime, with framework alias fallback."""
    runtime = (runtime_framework or "pytorch").strip().lower()
    ordered_keys = ("onnx", "torchscript", "tensorrt", "openvino", "pytorch")

    for key in ordered_keys:
        if key in runtime and key in loaders:
            return loaders[key](action_tracker)

    if default_key in loaders:
        return loaders[default_key](action_tracker)

    if "pytorch" in loaders:
        return loaders["pytorch"](action_tracker)

    if loaders:
        return next(iter(loaders.values()))(action_tracker)

    raise ValueError("No model loader was provided.")


@dataclass
class ModelLoaderRegistry:
    """Runtime-keyed model loader registry."""

    loaders: dict[str, ModelLoaderCallable] = field(default_factory=dict)
    default_key: str = "pytorch"

    def register(self, runtime_key: str, loader: ModelLoaderCallable) -> None:
        """Register a model loader for a runtime key."""
        self.loaders[runtime_key.strip().lower()] = loader

    def resolve(self, runtime_framework: str) -> ModelLoaderCallable:
        """Resolve a model loader for a runtime descriptor string."""
        runtime = (runtime_framework or "pytorch").strip().lower()
        for key in ("onnx", "torchscript", "tensorrt", "openvino", "pytorch"):
            if key in runtime and key in self.loaders:
                return self.loaders[key]
        if self.default_key in self.loaders:
            return self.loaders[self.default_key]
        if self.loaders:
            return next(iter(self.loaders.values()))
        raise ValueError("No model loader was registered.")

