"""Task-aware runtime registry for eval model loaders and runners."""

from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from .eval_model_loader import ModelLoaderCallable, load_model_with_runtime
from .exceptions import RuntimeNotSupportedError
from .inference_runners import RunnerCallable, get_inference_runner


if TYPE_CHECKING:
    from collections.abc import MutableMapping


def _norm(value: str, fallback: str) -> str:
    """Normalize registry keys with fallback and lowercase formatting.

    Args:
        value: Candidate key provided by caller.
        fallback: Default key when ``value`` is empty.

    Returns:
        Normalized lowercase key string.
    """
    return (value or fallback).strip().lower()


@dataclass
class RuntimeTaskRegistry:
    """Registry keyed by task type and runtime framework."""

    default_runtime: str = "pytorch"
    loader_map: MutableMapping[str, dict[str, ModelLoaderCallable]] = field(
        default_factory=lambda: defaultdict(dict)
    )
    runner_map: MutableMapping[str, dict[str, RunnerCallable]] = field(
        default_factory=lambda: defaultdict(dict)
    )

    def register_loader(self, task_type: str, runtime_key: str, loader: ModelLoaderCallable) -> None:
        """Register a model loader for a task/runtime pair.

        Args:
            task_type: Task namespace to register loader under.
            runtime_key: Runtime key associated with the loader.
            loader: Model loader callable for the task/runtime pair.

        Returns:
            None.
        """
        self.loader_map[_norm(task_type, "generic")][_norm(runtime_key, self.default_runtime)] = loader

    def register_runner(self, task_type: str, runtime_key: str, runner: RunnerCallable) -> None:
        """Register an inference runner for a task/runtime pair.

        Args:
            task_type: Task namespace to register runner under.
            runtime_key: Runtime key associated with the runner.
            runner: Inference runner callable for the task/runtime pair.

        Returns:
            None.
        """
        self.runner_map[_norm(task_type, "generic")][_norm(runtime_key, self.default_runtime)] = runner

    def resolve_loader(self, task_type: str, runtime_framework: str) -> ModelLoaderCallable:
        """Resolve a model loader callable for a task/runtime descriptor.

        Args:
            task_type: Task namespace used for loader lookup.
            runtime_framework: Runtime descriptor used for runtime matching.

        Returns:
            Loader callable configured for resolved task/runtime pair.
        """
        task = _norm(task_type, "generic")
        loaders = self.loader_map.get(task) or self.loader_map.get("generic")
        if not loaders:
            raise RuntimeNotSupportedError(
                f"No loaders registered for task='{task}'."
            )

        def _wrapped(action_tracker: Any) -> Any:
            """Load a model with runtime fallback for the current task.

            Args:
                action_tracker: Tracker-like object used by loader callbacks.

            Returns:
                Loaded model artifact for the requested runtime framework.
            """
            return load_model_with_runtime(
                action_tracker=action_tracker,
                runtime_framework=runtime_framework,
                loaders=loaders,
                default_key=self.default_runtime,
            )

        return _wrapped

    def resolve_runner(self, task_type: str, runtime_framework: str) -> RunnerCallable:
        """Resolve an inference runner callable for a task/runtime descriptor.

        Args:
            task_type: Task namespace used for runner lookup.
            runtime_framework: Runtime descriptor used for runtime matching.

        Returns:
            Runner callable configured for resolved task/runtime pair.
        """
        task = _norm(task_type, "generic")
        runners = self.runner_map.get(task) or self.runner_map.get("generic")
        if not runners:
            raise RuntimeNotSupportedError(
                f"No runners registered for task='{task}'."
            )
        return get_inference_runner(
            runtime_framework=runtime_framework,
            runners=runners,
            default_key=self.default_runtime,
        )

