"""Framework inference runner selection utilities."""

from __future__ import annotations

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


RunnerCallable = Callable[..., object]


def _normalize_runtime(runtime_framework: str) -> str:
    """Normalize a runtime descriptor to a canonical lowercase key.

    Args:
        runtime_framework: Runtime identifier from config or caller input.

    Returns:
        Lowercased runtime key, defaulting to ``pytorch`` when blank.
    """
    return (runtime_framework or "pytorch").strip().lower()


def get_inference_runner(
    runtime_framework: str,
    runners: dict[str, RunnerCallable],
    default_key: str = "pytorch",
) -> RunnerCallable:
    """Resolve a runtime framework to a registered inference callable.

    Args:
        runtime_framework: Runtime descriptor string from caller.
        runners: Mapping of runtime keys to runner callables.
        default_key: Default key used when runtime descriptor is ambiguous.

    Returns:
        Resolved inference runner callable.
    """
    runtime = _normalize_runtime(runtime_framework)
    keys_to_try: Iterable[str] = (
        "onnx",
        "torchscript",
        "tensorrt",
        "openvino",
        "pytorch",
    )

    for key in keys_to_try:
        if key in runtime and key in runners:
            return runners[key]

    if default_key in runners:
        return runners[default_key]

    if "pytorch" in runners:
        return runners["pytorch"]

    if runners:
        return next(iter(runners.values()))

    raise ValueError("No inference runners were provided.")


@dataclass
class InferenceRunnerRegistry:
    """Registry used by evaluators to select framework-specific runners."""

    runners: dict[str, RunnerCallable] = field(default_factory=dict)
    default_key: str = "pytorch"

    def register(self, runtime_key: str, runner: RunnerCallable) -> None:
        """Register an inference runner under a runtime key.

        Args:
            runtime_key: Runtime key used during runner resolution.
            runner: Runner callable to register.

        Returns:
            None.
        """
        self.runners[runtime_key.strip().lower()] = runner

    def resolve(self, runtime_framework: str) -> RunnerCallable:
        """Resolve the best runner for a runtime descriptor string.

        Args:
            runtime_framework: Runtime descriptor requested by caller.

        Returns:
            Resolved inference runner callable.
        """
        return get_inference_runner(runtime_framework, self.runners, self.default_key)

    def get(self, runtime_key: str) -> RunnerCallable | None:
        """Fetch a runner by exact runtime key when present.

        Args:
            runtime_key: Exact runtime key to look up.

        Returns:
            Registered runner callable, or ``None`` when key is missing.
        """
        return self.runners.get(runtime_key.strip().lower())

