"""Shared model-loading primitives used across multiple packages."""

from __future__ import annotations

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


if TYPE_CHECKING:
    from collections.abc import Callable


@dataclass(frozen=True)
class ModelLoadStrategies:
    """Strategy callbacks for checkpoint/predefined/scratch model loading.

    Args:
        from_checkpoint: Callback to load model from checkpoint path.
        from_predefined: Callback to load model from predefined model value.
        from_scratch: Optional callback to create model from scratch.
    """

    from_checkpoint: Callable[[str], Any]
    from_predefined: Callable[[str], Any]
    from_scratch: Callable[[], Any] | None = None


def load_model_with_strategies(action_tracker: Any, strategies: ModelLoadStrategies) -> Any:
    """Load model using checkpoint, predefined, then scratch fallback order.

    Args:
        action_tracker: Tracker-like object exposing ``get_checkpoint_path``.
        strategies: Strategy callbacks for each loading branch.

    Returns:
        Loaded model artifact from the first available strategy.
    """
    model_checkpoint, predefined_value = action_tracker.get_checkpoint_path()
    if model_checkpoint:
        return strategies.from_checkpoint(model_checkpoint)
    if predefined_value:
        return strategies.from_predefined(predefined_value)
    if strategies.from_scratch is not None:
        return strategies.from_scratch()
    raise ValueError("Neither checkpoint nor predefined model value is available.")


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 model using runtime descriptor with alias and fallback support.

    Args:
        action_tracker: Tracker-like object passed to selected loader.
        runtime_framework: Runtime descriptor string requested by caller.
        loaders: Mapping of runtime keys to loader callables.
        default_key: Default loader key used when runtime is ambiguous.

    Returns:
        Loaded model artifact from selected loader.
    """
    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.

    Args:
        loaders: Mapping from runtime key to loader callable.
        default_key: Default runtime key used when no direct match is found.
    """

    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.

        Args:
            runtime_key: Runtime key used during loader resolution.
            loader: Loader callable for the runtime key.

        Returns:
            None.
        """
        self.loaders[runtime_key.strip().lower()] = loader

    def resolve(self, runtime_framework: str) -> ModelLoaderCallable:
        """Resolve a model loader for a runtime descriptor string.

        Args:
            runtime_framework: Runtime descriptor requested by caller.

        Returns:
            Resolved model loader callable.
        """
        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.")
