from __future__ import annotations

from typing import TYPE_CHECKING, Any, TypeVar, cast

if TYPE_CHECKING:
    from collections.abc import Callable

    from kele.control.grounding_selector._rule_strategies.strategy_protocol import RuleSelectionStrategy
    from kele.control.grounding_selector._term_strategies.strategy_protocol import TermSelectionStrategy
else:
    RuleSelectionStrategy = Any
    TermSelectionStrategy = Any

StrategyT = TypeVar("StrategyT")


class StrategyRegistry[StrategyT]:
    """Registry for named strategy implementations within a single task."""

    def __init__(self, task_name: str) -> None:
        self._task_name = task_name
        self._registry: dict[str, type[StrategyT]] = {}

    @property
    def task_name(self) -> str:
        """Return the task name for this registry."""
        return self._task_name

    def register(self, name: str) -> Callable[[type[StrategyT]], type[StrategyT]]:
        """Return a decorator that registers a strategy class under ``name``."""
        def decorator(cls: type[StrategyT]) -> type[StrategyT]:
            self._registry[name] = cls
            return cls

        return decorator

    def __call__(self, name: str) -> Callable[[type[StrategyT]], type[StrategyT]]:
        """Allow the registry instance to be used as a decorator factory."""
        return self.register(name)

    def get(self, name: str) -> type[StrategyT]:
        """Return the registered strategy class by name."""
        try:
            return self._registry[name]
        except KeyError as err:
            raise KeyError(
                f"{self._task_name} strategy '{name}' is not registered. Available strategies: "
                f"{list(self._registry.keys())}"
            ) from err

    def available(self) -> list[str]:
        """Return the list of registered strategy names."""
        return list(self._registry.keys())


class RegistryHub:
    """Registry hub that exposes the supported task-specific registries."""

    def __init__(self) -> None:
        self._tasks: dict[str, StrategyRegistry[Any]] = {}

    def _task(self, name: str) -> StrategyRegistry[Any]:
        """Get or create a task registry by name for internal extension points."""
        if name not in self._tasks:
            self._tasks[name] = StrategyRegistry(name)
        return self._tasks[name]

    @property
    def rule_selector(self) -> StrategyRegistry[RuleSelectionStrategy]:
        """Registry for rule selection strategies."""
        return cast("StrategyRegistry[RuleSelectionStrategy]", self._task("rule_selector"))

    @property
    def term_selector(self) -> StrategyRegistry[TermSelectionStrategy]:
        """Registry for term selection strategies."""
        return cast("StrategyRegistry[TermSelectionStrategy]", self._task("term_selector"))


_registry = RegistryHub()


def register_rule_strategy(name: str) -> Callable[[type[RuleSelectionStrategy]], type[RuleSelectionStrategy]]:
    """Register a rule selection strategy class under ``name``."""
    return _registry.rule_selector.register(name)


def register_term_strategy(name: str) -> Callable[[type[TermSelectionStrategy]], type[TermSelectionStrategy]]:
    """Register a term selection strategy class under ``name``."""
    return _registry.term_selector.register(name)


def get_rule_strategy_class(name: str) -> type[RuleSelectionStrategy]:
    """Return the registered rule selection strategy class by name."""
    return _registry.rule_selector.get(name)


def get_term_strategy_class(name: str) -> type[TermSelectionStrategy]:
    """Return the registered term selection strategy class by name."""
    return _registry.term_selector.get(name)


register = _registry

__all__ = [
    "RegistryHub",
    "StrategyRegistry",
    "get_rule_strategy_class",
    "get_term_strategy_class",
    "register",
    "register_rule_strategy",
    "register_term_strategy",
]
