from __future__ import annotations
from typing import (
    Type,
    TypeVar,
    Literal,
    Generator,
    Iterator,
    overload,
    Any,
    ClassVar,
)
from types import GenericAlias

import sys
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, is_dataclass, fields, asdict, InitVar
import traceback


T = TypeVar("T")
ContextDataType = TypeVar("ContextDataType", bound="ContextData")
ContextType = TypeVar("ContextType", bound="Context")


class Context:
    _ctx_var: ContextVar["Context"] = ContextVar("Context")

    @classmethod
    def _decorate_context_type(cls, class_to_decorate):
        # Defaults:
        dataclass_args = {
            "frozen": True,
            "repr": True,
        }
        # Optimizations for python 3.10+
        if sys.version_info >= (3, 10):
            dataclass_args["slots"] = True
            dataclass_args["kw_only"] = True
        decorated = dataclass(class_to_decorate, **dataclass_args)
        return decorated

    @classmethod
    def context_type(cls, class_to_decorate: Type[T]) -> Type[T]:
        if not issubclass(class_to_decorate, ContextData):
            raise TypeError("You must inherit ContextData!")
        for name in class_to_decorate.__annotations__.keys():
            if name in ("_depth", "_on_entering", "_on_exiting", "_on_exited"):
                raise TypeError(
                    f"You cannot declare the ContextData subclass {class_to_decorate} with a field {name!r}: it is a reserved keyword."
                )

        decorated_class = cls._decorate_context_type(class_to_decorate)
        return decorated_class

    @overload
    @classmethod
    def current(
        cls: Type[ContextType], raises: Literal[False]
    ) -> ContextType | None: ...

    @overload
    @classmethod
    def current(cls: Type[ContextType], raises: Literal[True]) -> ContextType: ...

    @overload
    @classmethod
    def current(cls: Type[ContextType]) -> ContextType: ...

    @classmethod
    def current(cls: Type[ContextType], raises: bool = True) -> ContextType | None:
        ctx = cls._ctx_var.get(None)
        if ctx is None:
            if not raises:
                return None
            raise ValueError(f"No context is currently active.")
        if not isinstance(ctx, cls):
            if not raises:
                return None
            raise TypeError(f"The current context is a {ctx.__class__}, not a {cls}.")
        return ctx

    @classmethod
    def _augment(
        cls,
        parent: ContextData | None,
        SubType: Type[ContextDataType],
        kwargs: dict[str, Any],
        skip_extra_fields: bool,
    ) -> ContextDataType:
        sub_kwargs = {}
        if parent is not None:
            parent_fields = fields(parent)  # type: ignore - we know it is a dataclass
            sub_kwargs = {f.name: getattr(parent, f.name) for f in parent_fields}
        sub_kwargs.update(**kwargs)
        if skip_extra_fields:
            filtered = {}
            for field in fields(SubType):  # type: ignore - we know it is a dataclass
                filtered[field.name] = sub_kwargs[field.name]
            sub_kwargs = filtered
        depth = sub_kwargs.get("_depth", 0)
        sub_kwargs["_depth"] = depth + 1
        # TODO: should we validate all field type here?
        # NB: I tried to use pydantic.dataclasses.dataclass for that but it's not that straightforward
        # If we decide that ContextData must validate their fields, we should probably drop dataclasses and use pydantic models...
        try:
            sub = SubType(**sub_kwargs)
        except TypeError as err:
            raise TypeError(
                f"While augmenting a {parent.__class__.__name__} to a {SubType.__name__}: {err}"
            )
        return sub

    @classmethod
    @contextmanager
    def _get(
        cls,
        ContextDataType: Type[ContextDataType],
        kwargs: Any,
        skip_extra_fields: bool,
    ) -> Generator[ContextDataType, None, None]:
        assert cls is Context or is_dataclass(ContextDataType)
        parent: ContextDataType = cls.current(raises=False)  # type: ignore - we know it is a ContextData
        sub = cls._augment(
            parent, ContextDataType, kwargs, skip_extra_fields=skip_extra_fields
        )
        token = cls._ctx_var.set(sub)
        with sub._on_entering():
            try:
                yield sub
            finally:
                sub._on_exiting()
                cls._ctx_var.reset(token)
                sub._on_exited()

    @classmethod
    @contextmanager
    def get(
        cls, ContextDataType: Type[ContextDataType], **kwargs: Any
    ) -> Generator[ContextDataType, None, None]:
        """
        Create a new context of type ContextDataType.
        The current context must be a subclass of ContextDataType (or ContextDataType must support the
        fields of the current context)
        """
        with cls._get(ContextDataType, kwargs, skip_extra_fields=False) as ctx:
            yield ctx

    @classmethod
    @contextmanager
    def get_as(
        cls, ContextDataType: Type[ContextDataType], **kwargs: Any
    ) -> Generator[ContextDataType, None, None]:
        """
        Create a new context of type ContextDataType.
        Fields of the current context are dropped if not supported
        by ContextDataType.
        """
        with cls._get(ContextDataType, kwargs, skip_extra_fields=True) as ctx:
            yield ctx


@Context._decorate_context_type
class ContextData(Context):
    _depth: int

    @contextmanager
    def _on_entering(self):
        """
        Called right before this context becomes current.
        """
        yield

    def _on_exiting(self):
        """
        Called right before this context exits.
        """
        pass

    def _on_exited(self):
        """
        Called right after this context exited.
        """
        pass
