from .context import Context, ContextData
from .loggable_context import LoggableContext


# 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


# import rich
# import rich.console
# import rich.panel
# import rich.table
# import rich.columns


# T = TypeVar("T")


# 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


# ContextType = TypeVar("ContextType", bound=Context)


# @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


# class RichLogger:
#     def __init__(self, record: bool = False, indent_width=4) -> None:
#         self._console = rich.console.Console(record=record)
#         self._indent = 0
#         self._indent_width = max(2, indent_width)

#     def export_text(self) -> str:
#         return self._console.export_text()

#     def _indented(self, r, lc="│ "):
#         sep = "│"
#         i = f"{{:>{self._indent_width}}}".format(sep)

#         ilc = f"??{{:>{self._indent_width}}}".format(lc)
#         ilc = ((self._indent_width - 1) * " ") + lc
#         col = rich.columns.Columns(
#             [(self._indent // 2) * i + ilc, r],
#             padding=(0, 0),
#             # column_first=True,
#             align="right",
#         )
#         return col

#     def message(self, message: rich.console.RenderableType):
#         r = self._indented(message, lc="├─ ")
#         self._console.print(r)

#     def log_dict(
#         self,
#         d: dict[rich.console.RenderableType, rich.console.RenderableType],
#         title: str | None = None,
#     ):
#         table = rich.table.Table(title=title)
#         for k, v in d.items():
#             table.add_row(str(k), repr(v))
#         self._console.print(self._indented(table))

#     def _section_header(self, title: str):
#         p = rich.panel.Panel(title, expand=False)
#         return p

#     def _section_footer(self):
#         return ""

#     @contextmanager
#     def group(self, title: str):
#         self.message(title)
#         self._indent += 2
#         try:
#             yield
#         finally:
#             # self._indent -= 1
#             self._console.print(self._indented("", lc=f"╰🠆 ✓"))
#             self._indent -= 2

#     @contextmanager
#     def section(self, title: str):
#         w = len(title)
#         i = (self._indent_width - 2) * "─"
#         self._console.print(self._indented("", lc="│╭─" + w * "─" + "─╮"))
#         self._console.print(self._indented("", lc="├┤ " + title + " │"))
#         self._console.print(
#             self._indented("", lc=f"│╰{i}┬" + (w - self._indent_width + 2) * "─" + "─╯")
#         )
#         # self._console.print(rich.padding.Padding.indent(p, (indent * 4) - 2))
#         self._indent += 2
#         try:
#             yield
#         except Exception as err:

#             if footer := self._section_footer():
#                 self._console.print(self._indented("", lc=f"├🠆 {footer}"))

#             self._console.print(self._indented("[red]✘ Error:[/red]", lc=f"├🠆 "))
#             tb = traceback.format_exc()
#             self._console.print(rich.panel.Panel(tb, style="red"))
#             if str(err):
#                 self._console.print(self._indented(f"[red]{err}[/red]", lc=f"├🠆 "))

#         else:
#             if footer := self._section_footer():
#                 self._console.print(self._indented("", lc=f"├🠆 {footer}"))
#             msg = f"[green]🗹[/green] Done {title}"
#             self._console.print(self._indented("", lc=f"╰🠆 {msg}"))
#         finally:
#             self._indent -= 2
#             # self._console.print(self._indented("", lc="├─╯"))


# @Context._decorate_context_type
# class LoggableContext(ContextData):
#     @classmethod
#     def create_logger(cls, record: bool = False) -> RichLogger:
#         return RichLogger(record=record)

#     logger: RichLogger | None

#     def _log_section_summary(self) -> str:
#         return ""

#     @contextmanager
#     def _on_entering(self):
#         if self.logger is not None:
#             with self.logger.section(self._log_section_summary()):
#                 yield
#         return

#     def log(self, message: str):
#         if self.logger is None:
#             return
#         self.logger.message(message)

#     @contextmanager
#     def log_group(self, title: str):
#         if self.logger is None:
#             return
#         with self.logger.group(title):
#             yield

#     @contextmanager
#     def log_section(self, title: str):
#         if self.logger is None:
#             return
#         with self.logger.section(title):
#             yield


# ContextDataType = TypeVar("ContextDataType", bound=ContextData)
