"""Framework routing helpers for multi-runtime deploy flows."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from .contracts import ModelBundle, RouteRegistry, RuntimeRoute
from .runtime_resolver import normalize_runtime, resolve_runtime_framework


if TYPE_CHECKING:
    from collections.abc import Callable, Mapping


class FrameworkRouter:
    """Router that dispatches load/predict handlers by runtime key."""

    def __init__(
        self,
        routes: RouteRegistry,
        *,
        default_runtime: str = "pytorch",
        fallback_runtime: str | None = None,
    ) -> None:
        """Initialize runtime routes and fallback selection behavior.

        Args:
            routes: Mapping of runtime keys to route handler objects.
            default_runtime: Runtime key used when none is requested.
            fallback_runtime: Optional runtime used when primary load fails.

        Returns:
            None.
        """
        if not routes:
            raise ValueError("FrameworkRouter requires at least one runtime route.")
        self._routes: dict[str, RuntimeRoute] = {
            normalize_runtime(runtime): route for runtime, route in routes.items()
        }
        self.default_runtime = normalize_runtime(default_runtime)
        self.fallback_runtime = (
            normalize_runtime(fallback_runtime) if fallback_runtime is not None else None
        )

    def supported_runtimes(self) -> tuple[str, ...]:
        """Return registered runtime keys.

        Args:
            None.

        Returns:
            Sorted tuple of normalized runtime keys.
        """
        return tuple(sorted(self._routes.keys()))

    def _resolve_route(self, runtime_framework: str | None) -> tuple[str, RuntimeRoute]:
        """Resolve runtime name to an effective route with fallback support.

        Args:
            runtime_framework: Requested runtime key or ``None``.

        Returns:
            Tuple of resolved runtime key and its route handler.
        """
        runtime = normalize_runtime(runtime_framework or self.default_runtime)
        route = self._routes.get(runtime)
        if route is not None:
            return runtime, route
        if self.fallback_runtime and self.fallback_runtime in self._routes:
            return self.fallback_runtime, self._routes[self.fallback_runtime]
        raise KeyError(
            f"Unsupported runtime '{runtime}'. Supported: {', '.join(self.supported_runtimes())}"
        )

    def load_model(self, action_tracker: Any, runtime_framework: str | None = None) -> ModelBundle:
        """Load model via selected runtime route with optional fallback.

        Args:
            action_tracker: Tracker-like object forwarded to load handler.
            runtime_framework: Requested runtime key or ``None``.

        Returns:
            Loaded model bundle containing runtime, model, and metadata.
        """
        runtime, route = self._resolve_route(runtime_framework)
        try:
            model = route.loader(action_tracker)
            return ModelBundle(framework=runtime, model=model)
        except Exception as error:
            if self.fallback_runtime and runtime != self.fallback_runtime:
                fallback_route = self._routes.get(self.fallback_runtime)
                if fallback_route is not None:
                    model = fallback_route.loader(action_tracker)
                    return ModelBundle(
                        framework=self.fallback_runtime,
                        model=model,
                        metadata={
                            "requested_runtime": runtime,
                            "fallback_runtime": self.fallback_runtime,
                            "load_error": str(error),
                        },
                    )
            raise

    def predict(self, bundle: ModelBundle, input_bytes: bytes) -> Any:
        """Run runtime-specific prediction for a loaded model bundle.

        Args:
            bundle: Model bundle produced by ``load_model``.
            input_bytes: Serialized request payload passed to predictor.

        Returns:
            Runtime-specific prediction result.
        """
        runtime = normalize_runtime(bundle.framework)
        route = self._routes.get(runtime)
        if route is None:
            raise KeyError(
                f"No predict route registered for runtime '{runtime}'. "
                f"Supported: {', '.join(self.supported_runtimes())}"
            )
        return route.predictor(bundle.model, input_bytes)


def build_router_functions(
    router: FrameworkRouter,
    *,
    runtime_getter: Callable[[Any], str] | None = None,
) -> tuple[Callable[[Any], ModelBundle], Callable[[ModelBundle, bytes], Any]]:
    """Build deploy-compatible ``(load_model, predict)`` callables.

    Args:
        router: Router that resolves runtime-specific handlers.
        runtime_getter: Optional callable to derive runtime from tracker.

    Returns:
        Tuple of ``load_model`` and ``predict`` callables for server wiring.
    """
    runtime_resolver = runtime_getter or (lambda tracker: resolve_runtime_framework(tracker))

    def load_model(action_tracker: Any) -> ModelBundle:
        """Load a model bundle using runtime inferred from tracker state.

        Args:
            action_tracker: Tracker object used for runtime resolution.

        Returns:
            Loaded model bundle for the resolved runtime.
        """
        runtime = runtime_resolver(action_tracker)
        return router.load_model(action_tracker, runtime)

    def predict(model_bundle: ModelBundle, input_bytes: bytes) -> Any:
        """Run prediction through router for an existing model bundle.

        Args:
            model_bundle: Loaded model bundle returned by ``load_model``.
            input_bytes: Serialized request payload for prediction.

        Returns:
            Runtime-specific prediction output.
        """
        return router.predict(model_bundle, input_bytes)

    return load_model, predict


def make_routes(route_map: Mapping[str, tuple[Callable[[Any], Any], Callable[[Any, bytes], Any]]]) -> dict[str, RuntimeRoute]:
    """Build runtime route objects from a mapping of callables.

    Args:
        route_map: Mapping from runtime key to ``(loader, predictor)`` tuple.

    Returns:
        Dictionary of runtime keys mapped to ``RuntimeRoute`` instances.
    """
    return {
        runtime: RuntimeRoute(loader=handlers[0], predictor=handlers[1])
        for runtime, handlers in route_map.items()
    }
