from __future__ import annotations

from typing import TYPE_CHECKING, ParamSpec, Self, TypeVar, override

from pydantic import BaseModel, ConfigDict, PrivateAttr

if TYPE_CHECKING:
  from collections.abc import Callable

__all__ = ["ConfigurableIndicator", "MakeableModel"]

P = ParamSpec("P")
R = TypeVar("R")


class MakeableModel[R](BaseModel):
  """Base class for generated configuration models.

  Generated config models inherit from this class and override the make()
  method to create instances or bound callables with configured hyperparameters.

  For classes/dataclasses: make() returns the instance directly.
  For functions: make() returns a callable with hyperparameters bound.
  """

  model_config = ConfigDict()

  _instance: R | None = PrivateAttr(default=None)

  @classmethod
  def from_instance(cls, instance: R) -> Self:
    """Create a Config that wraps an existing instance.

    Args:
      instance: The pre-existing instance to wrap.

    Returns:
      A Config object that will return the wrapped instance when make() is called.
    """
    model = cls.model_construct()
    model._instance = instance
    return model

  def make(self) -> R:
    """Creates an instance or callable from this configuration.

    Returns:
      For classes: The instance with configured hyperparameters
      For functions: A callable with hyperparameters pre-bound

    Raises:
      NotImplementedError: This base implementation should be overridden
    """
    if self._instance is not None:
      return self._instance

    raise NotImplementedError(
      "This method should be overridden in the dynamic subclass."
    )

  @override
  def __str__(self) -> str:
    return repr(self)


class ConfigurableIndicator[**P, R]:
  """Wrapper for @configurable decorated functions.

  Holds the original function and its generated Config class. The function
  remains directly callable while exposing the Config class for configuration.

  Attributes:
    Config: Generated Pydantic model class for this function's hyperparameters
  """

  Config: type[MakeableModel[R]]
  Type: type

  def __init__(self, fn: Callable[P, R], config_cls: type[MakeableModel[R]]) -> None:
    super().__init__()
    self._fn = fn
    self.Config = config_cls

    # Create a proxy type to allow 'my_func.Type' syntax for type hinting
    # This fools the extraction logic into treating it like a configurable class
    class ProxyType:
      Config = config_cls

    # Name it nicely for debugging
    ProxyType.__name__ = f"{fn.__name__}Type"
    ProxyType.__qualname__ = f"{fn.__qualname__}.Type"
    self.Type = ProxyType

  def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
    """Call the wrapped function directly."""
    return self._fn(*args, **kwargs)

  def __get__(
    self, obj: object, objtype: type | None = None
  ) -> ConfigurableIndicator[P, R]:
    """Support method binding via descriptor protocol.

    Returns:
      Unbound ConfigurableIndicator if accessed from class,
      bound ConfigurableIndicator if accessed from instance
    """
    if obj is None:
      return self
    bound_fn = self._fn.__get__(obj, objtype)
    return ConfigurableIndicator(bound_fn, self.Config)
