from __future__ import annotations

import contextvars
import inspect
import textwrap
from collections import defaultdict
from datetime import timedelta
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Dict,
    Iterator,
    List,
    Optional,
    Tuple,
    Type,
    cast,
    Set,
    Protocol,
    Mapping,
)

from typing_extensions import TypeGuard

from chalk._lsp._class_finder import get_class_ast
from chalk.features.feature_wrapper import FeatureWrapper, unwrap_feature
from chalk.utils import notebook
from chalk.utils.collections import ensure_tuple, get_unique_item
from chalk.utils.source_parsing import should_skip_source_code_parsing

if TYPE_CHECKING:
    from chalk.features.feature_field import Feature
    from chalk._lsp.error_builder import FeatureClassErrorBuilder
    from chalk.features.underscore_features import NamedUnderscoreExpr

__all__ = ["Features", "FeaturesMeta", "FeatureSetBase", "is_features_cls", "is_feature_set_class"]


class ClassSource:
    __slots__ = ("filename", "class_source", "tree", "dedent_source")

    def __init__(self, c: Any):
        self.filename = None
        self.class_source = None
        self.tree = None
        super().__init__()
        if not should_skip_source_code_parsing():
            try:
                self.filename = inspect.getfile(c)
            except:
                self.filename = None

            try:
                if notebook.is_defined_in_module(c):
                    self.class_source = inspect.getsource(c)
                # TODO (rkargon) Try to get source from notebook cells if possible
                else:
                    self.class_source = None
            except:
                self.class_source = None
            self.dedent_source = self.class_source and textwrap.dedent(self.class_source)
            try:
                self.tree = get_class_ast(c)
            except:
                self.tree = None


class FeaturesMeta(type):
    """Metaclass for classes decorated with ``@features``.

    This metaclass allows for:

    1.  Classes annotated with @features to pass the
        ``isinstance(x, Features)`` and ``issubclass(X, Features)`` checks.
    2.  ``Features[Feature1, Feature2]`` annotations to return subclasses of Features, so they can be used as proper type annotations.
    """

    _features_lazy: Callable[[], List[Feature]]  # pyright: ignore[reportUninitializedInstanceVariable]

    @property
    def features(self) -> List[Feature]:
        return self._features_lazy()

    def __subclasscheck__(self, subcls: type) -> bool:
        if not isinstance(subcls, type):  # pyright: ignore[reportUnnecessaryIsInstance]
            raise TypeError(f"{subcls} is not a type")
        if hasattr(subcls, "__is_features__"):
            return getattr(subcls, "__is_features__")
        return False

    def __instancecheck__(self, instance: object) -> bool:
        return self.__subclasscheck__(type(instance))

    def __getitem__(cls, item: Any) -> Any:
        # This lets us subscript by the features
        # Annotating the return type as Type[Any] as instances of @features classes are
        # not recognized as being subclasses of Features by the type checker (even though at runtime they would be)
        from chalk.features.dataframe import DataFrame
        from chalk.features.feature_field import Feature

        is_exploded = False

        # Typing the `__getitem__` as any, since the @features members are typed as the underlying data structure
        # But, they should always be features or a tuple of features
        if isinstance(item, type) and issubclass(item, Features):
            is_exploded = True
            item = [
                f
                for f in item.features
                if not f.is_autogenerated and not f.is_has_one and not f.is_has_many and not f.is_windowed
            ]

        def f():
            item_ = ensure_tuple(item)
            item_ = tuple(unwrap_feature(x) if isinstance(x, FeatureWrapper) else x for x in item_)
            for x in item_:
                if isinstance(x, str):
                    raise TypeError(
                        (
                            f'String features like {cls.__name__}["{x}"] are unsupported. '
                            f'Instead, replace with "{cls.__name__}[{x}]"'
                        )
                    )

                if isinstance(x, Feature) or (isinstance(x, type) and issubclass(x, DataFrame)):  # pyright: ignore[reportUnnecessaryIsInstance]
                    continue

                raise TypeError(f"Invalid feature {x} of type '{type(x).__name__}'")
            cls_ = cast(Type[Features], cls)

            new_features = list([*getattr(cls_, "features", ()), *item_])
            if len(new_features) == 0:
                raise ValueError("Cannot create empty Features[] class")
            return new_features

        class SubFeatures(cls):  # pyright: ignore
            __is_exploded__ = is_exploded
            _features_lazy = f

        return SubFeatures

    def __repr__(cls) -> str:
        if not hasattr(cls, "features"):
            return "Features"
        return f"Features[{', '.join(str(f) for f in cls.features)}]"

    def __call__(cls, *args: object, **kwargs: object):
        raise RuntimeError("Instances features cls should never be directly created. Instead, use Features[User.id]")

    def __sub__(self, other: type[Features] | FeatureWrapper):
        if isinstance(other, type) and issubclass(other, Features):  # pyright: ignore[reportUnnecessaryIsInstance]
            other_features = set(other.features)
            return FeaturesImpl[tuple(f for f in self.features if f not in other_features)]
        if isinstance(other, FeatureWrapper):  # pyright: ignore[reportUnnecessaryIsInstance]
            return FeaturesImpl[tuple(f for f in self.features if f != unwrap_feature(other))]
        raise TypeError(f"Invalid type {type(other)} for subtraction with Features[...]")

    @property
    def namespace(cls):
        cls = cast(Type[Features], cls)
        return get_unique_item((x.namespace for x in cls.features), name=f"{cls.__name__} feature namespaces")


class FeaturesImpl(metaclass=FeaturesMeta):
    """Features base class.

    This class is never instantiated or directly inherited. However, classes
    annotated with @features can be thought of as inheriting from this class.
    It can be used with `isinstance` and `issubclass` checks, as well as for
    typing.
    """

    # Internally, the Features class is instantiated when results come through, and
    # results are bound to instances of this class via attributes
    __chalk_error_builder__: ClassVar[FeatureClassErrorBuilder]
    __chalk_is_singleton__: ClassVar[bool]
    __chalk_etl_offline_to_online__: ClassVar[bool]
    __chalk_max_staleness__: ClassVar[timedelta]
    __chalk_namespace__: ClassVar[str]
    __chalk_primary__: ClassVar[Feature | None] = None  # The primary key feature
    __chalk_owner__: ClassVar[Optional[str]]
    __chalk_tags__: ClassVar[List[str]]
    __chalk_ts__: ClassVar[Optional[Feature]]  # The timestamp feature
    features: ClassVar[List[Feature]]
    __chalk_features_raw__: ClassVar[List[Feature]]
    __chalk_source_info__: "ClassVar[ClassSource]"
    __chalk_is_loaded_from_notebook__: ClassVar[bool] = False
    """True if loaded on a branch server from a pickle"""
    __chalk_materialized_windows__: ClassVar[List[Feature]]
    __chalk_group_by_materialized_windows__: ClassVar[List[Feature]]
    __chalk_expression_windows__: ClassVar[List[Feature]]
    __chalk_notebook_feature_expressions__: ClassVar[Dict[str, NamedUnderscoreExpr]]

    __is_exploded__: ClassVar[bool] = False
    __is_features__: ClassVar[bool] = True

    # When constructing results, this class is instantiated

    def __iter__(self) -> Iterator[Tuple[str, Any]]:
        """Iterating over features yields tuples of (fqn, value) for all scalarish feature values."""
        raise NotImplementedError

    def __len__(self) -> int:
        """The number of features that are set."""
        raise NotImplementedError

    def items(self) -> Iterator[Tuple[str, Any]]:
        """Iterating over features yields tuples of (fqn, value) for all feature values."""
        raise NotImplementedError


if TYPE_CHECKING:

    class FeaturesProtocol(FeaturesImpl):
        def __class_getitem__(cls, item: Any) -> "FeaturesProtocol": ...

    Features = FeaturesProtocol
else:
    Features = FeaturesImpl


class FeatureRegistryProtocol(Protocol):
    def get_feature_sets(self) -> Mapping[str, Type[Features]]: ...

    def add_feature_set(self, updated_class: Type[Features]): ...

    def get_singletons(self) -> Mapping[str, Type[Features]]: ...


class FeatureSetBase:
    """Registry containing all @features classes."""

    registry: ClassVar[Dict[str, Type[Features]]] = {}
    """Mapping of fqn to Features cls"""

    __chalk_singletons__: ClassVar[Dict[str, Type[Features]]] = {}
    """Mapping of fqn to singleton Features cls. These are also included in `.registry`."""

    __chalk_notebook_defined_feature_fields__: ClassVar[Dict[str, Set[str]]] = defaultdict(set)
    """
    (Mapping of namespace to a set of unversioned attribute names).
    This tracks which feature fields have been added to a feature class in a notebook, after it's been defined.
    """

    hook: ClassVar[Optional[Callable[[Type[Features]], None]]] = None

    def __init__(self) -> None:
        super().__init__()
        raise RuntimeError("FeatureSetBase should never be instantiated")


class FeatureSetBaseWrapper(FeatureRegistryProtocol):
    """
    Wrapper around the global FeatureSetBase
    """

    def __init__(self):
        super().__init__()

    def get_feature_sets(self) -> Mapping[str, Type[Features]]:
        return FeatureSetBase.registry

    def add_feature_set(self, updated_class: Type[Features]):
        namespace = updated_class.namespace
        previous_features_class = FeatureSetBase.registry.get(namespace)
        if (
            notebook.is_notebook()
            and previous_features_class is not None
            and notebook.is_defined_in_module(previous_features_class)
        ):
            # Not generating an LSP here because we're in a notebook anyway
            # TODO: See if we can pretty-print lsp errors in notebooks, at which point we can generate one that points to the old feature class
            raise ValueError(
                f"Cannot re-define feature class '{previous_features_class.__name__}' in a notebook: it was previously defined in '{previous_features_class.__module__}'."
            )

        FeatureSetBase.registry[updated_class.__chalk_namespace__] = updated_class

        if updated_class.__chalk_is_singleton__:
            FeatureSetBase.__chalk_singletons__[updated_class.__chalk_namespace__] = updated_class

        if FeatureSetBase.hook is not None:
            try:
                FeatureSetBase.hook(updated_class)
            except:
                # If we were unsuccessful, restore the previous features class, and call the hook again
                if previous_features_class is not None:
                    FeatureSetBase.registry[updated_class.__chalk_namespace__] = previous_features_class
                    if updated_class.__chalk_is_singleton__:
                        FeatureSetBase.__chalk_singletons__[updated_class.__chalk_namespace__] = previous_features_class
                    FeatureSetBase.hook(previous_features_class)
                # Now raise the original exception
                raise

        ...

    def get_singletons(self) -> Mapping[str, Type[Features]]:
        return FeatureSetBase.__chalk_singletons__


_DEFAULT_FEATURE_REGISTRY: FeatureRegistryProtocol = FeatureSetBaseWrapper()
CURRENT_FEATURE_REGISTRY: contextvars.ContextVar[FeatureRegistryProtocol] = contextvars.ContextVar(
    "CURRENT_FEATURE_REGISTRY", default=_DEFAULT_FEATURE_REGISTRY
)


def is_features_cls(maybe_features: Any) -> TypeGuard[Type[Features]]:
    """BEWARE: this returns true for FeatureClass but also Features[...]"""
    return isinstance(maybe_features, type) and issubclass(maybe_features, Features)


def is_feature_set_class(maybe_features: Any) -> TypeGuard[Type[FeaturesImpl]]:
    return is_features_cls(maybe_features) and getattr(maybe_features, "__chalk_feature_set__", False)
