"""This module contains the Tremors Logger and Collector definition.

A Tremors logger can be used as a drop-in replacement for a standard
:class:`~logging.Logger`. Collectors may be attached to a logger. Each time
a message is logged, the collectors run if they are enabled for the message
level. When a collector runs, it may update its state. A logger adds the
current state of all of its collectors to any :class:`~logging.LogRecord`
that it produces.

A logger may also be used as a context manager. When used this way, loggers
maintain a hierarchy. Each logger has a parent logger, unless it is the root
logger. The root logger assigns a :attr:`~Logger.group_id` that is shared
by all loggers in the hierarchy. Each logger has a :attr:`~Logger.path` from
its root logger to itself. Hierarchies may be nested by explicitly creating
a new root logger within an existing hierarchy.  The loggers in the nested
hierarchy will have a different group ID than the loggers in the containing
hierarchy. The paths in the nested hierarchy will also start from the new
root logger.

When a logger context is entered, or exited, a message is automatically
logged at a specified level, and collectors will run if they are enabled
for the message level. This allows us to define collectors that measure
lifecycle information, such as the duration of a context, or how much memory
was allocated at the beginning, and end of a context.

Each Tremors logger uses an underlying standard logger that may
be specified. The underlying logger can be configured in the normal
fashion. Underlying loggers may be shared by Tremors loggers. In fact,
it is common for all Tremors loggers to use the standard root logger. The
Tremors logger simply adds the states of its collectors to the records
produced by its standard logger. These augmented records can then be
processed by any :class:`~logging.Filter`, :class:`~logging.Formatter`,
or :class:`~logging.Handler` that knows about the extra record attributes.

A collector may be bundled with a formatter that can extract the collector
state from a record, and format it. This formatter can be used in a
filter, formatter, or handler that has been configured for the underling
logger. Tremors comes with many useful collector bundles. But you can also
define custom collectors, and bundles.
"""

from __future__ import annotations

import abc
import contextvars
import functools
import itertools
import logging
import uuid
from collections.abc import Callable, Iterable, Mapping, MutableMapping
from typing import TYPE_CHECKING, Any, NamedTuple, Self

if TYPE_CHECKING:
    from types import TracebackType

_current_logger: contextvars.ContextVar[Logger | None] = contextvars.ContextVar(
    "_current_logger", default=None
)


class LogItem(NamedTuple):
    """The properties of a logged message."""

    logger: Logger | None
    level: int
    msg: object
    args: tuple[object, ...] | Mapping[str, object] | None
    exc_info: (
        bool
        | tuple[type[BaseException], BaseException, TracebackType | None]
        | tuple[None, None, None]
        | BaseException
        | None
    )
    stack_info: bool
    stacklevel: int | None
    extra: Mapping[str, object] | None
    record: logging.LogRecord | None


class Collector[TState](NamedTuple):
    """The collector specification.

    Attributes:
        id: A unique identifier for the collector.
        name: The name of the collector. When the collector runs
            for a logged message, the collector's state is added to the
            :class:`~logging.LogRecord`, and may be retrieved using this name.
        level: The minimum level a logged message must be for the collector
            to run.
        inherit: If True, the collector will also be added to descendant
            loggers, provided its name does not conflict that of any of the
            other collectors on the descendant loggers.
        state: The initial state of the collector. Each time the collector
            runs, it may update this state.
        collect: When the collector runs, this function is called with the
            current state, and the LogItem. It returns the new state.
    """

    id: uuid.UUID
    name: str
    level: int
    inherit: bool
    state: TState
    collect: Callable[[TState, LogItem], TState]


class CollectorFactory[TState](metaclass=abc.ABCMeta):
    """A factory that creates a collector."""

    @abc.abstractmethod
    def create(self) -> Collector[TState]:
        """Create a collector."""
        raise NotImplementedError


def _collectors_reducer[T](
    acc: tuple[MutableMapping[str, Collector[T]], MutableMapping[str, T]],
    curr: tuple[LogItem, Collector[T]],
) -> tuple[MutableMapping[str, Collector[T]], MutableMapping[str, T]]:
    acc_collectors, acc_extra = acc
    log_item, curr_c = curr
    name = curr_c.name
    if curr_c.level == logging.NOTSET or log_item.level < curr_c.level:
        acc_collectors[name] = curr_c
        acc_extra[name] = curr_c.state
        return acc_collectors, acc_extra
    state = curr_c.collect(curr_c.state, log_item)
    acc_collectors[name] = Collector(
        curr_c.id, name, curr_c.level, curr_c.inherit, state, curr_c.collect
    )
    acc_extra[name] = state
    return acc_collectors, acc_extra


def _names_and_ids_reducer(
    acc: tuple[set[str], set[uuid.UUID]], curr: str | uuid.UUID
) -> tuple[set[str], set[uuid.UUID]]:
    names, ids = acc
    match curr:
        case str():
            names.add(curr)
        case uuid.UUID():
            ids.add(curr)
    return acc


EXTRA_KEY = "tremors"
ENTER_KEY = "__enter__"
EXIT_KEY = "__exit__"


class Logger(logging.LoggerAdapter[logging.Logger]):
    """The Tremors logger.

    Args:
        *collectors: Collectors or CollectorFactories. The collectors, or
            resultant collectors will be attached to the logger, and will be run
            for every logged message for which the collector is enabled.
        name: The name of the logger. This name is logged in the default
            entered, and exited messages. This name is used to generate the
            path, but may be altered in the path to be unique in the hierarchy.
        logger_name: The name of the underlying :class:`~logging.Logger`
            that will be used to log messages. If None, the root logger will
            be used.
        ctx_level: The level at which entered and exited messages will
            be logged.
        enter_msg: If a non-empty string, this will be the logged entered
            message. If truthy, the default entered message will be logged.
            Otherwise, no entered message will be logged.
        exit_msg: If a non-empty string, this will be the logged exited
            message. If truthy, the default exited message will be logged.
            Otherwise, no exited message will be logged.
        is_root: If True, a root logger with no parent that starts a new
            hierarchy will be created.
    """

    def __init__(  # noqa: PLR0913
        self,
        *collectors: Collector[Any] | CollectorFactory[Any],
        name: str,
        logger_name: str | None = None,
        ctx_level: int = logging.INFO,
        enter_msg: str | bool = True,
        exit_msg: str | bool = True,
        is_root: bool = False,
    ) -> None:
        """Initialize the logger."""
        super().__init__(logging.getLogger(logger_name))
        self._name = name
        self._collectors: MutableMapping[str, Collector[Any]] = {
            (coll := c_or_fact.create()).name
            if isinstance(c_or_fact, CollectorFactory)
            else (coll := c_or_fact).name: coll
            for c_or_fact in collectors
        }
        self._entered = 0
        self._ctx_level = ctx_level
        self._enter_msg = enter_msg
        self._exit_msg = exit_msg
        self._cv_token: contextvars.Token[Logger | None] | None = None
        self._is_root = is_root
        self._parent: Logger | None = None
        self._group_id: uuid.UUID | None = None
        self._path: tuple[str, ...] | None = None
        self._path_registry: dict[str, int] = {}

    def __enter__(self) -> Self:
        """Enter the context.

        The logger will only be added to the hierarchy, and log the entered
        message if the logger is not currently entered, i.e., it has not
        been entered before, or the number of times it has been entered,
        and exited are balanced.

        If the logger is added to the hierarchy, an entered message will be
        logged at the ``ctx_level`` the logger was initialized with. These
        messages may be suppressed by specifying a sufficiently low
        level. These messages may also be filtered out. The message text
        will be ``"entered: %s" % name`` where ``name`` is the name the
        logger was initialized with.
        """
        self._entered += 1
        if self._entered != 1:
            return self
        if not self._is_root:
            self._parent = _current_logger.get()
        if self._parent:
            self._group_id = self._parent.group_id
            self._path = self._parent.register_path(self)
            for name, collector in self._parent.collectors.items():
                if collector.inherit and name not in self._collectors:
                    self._collectors[name] = collector
        else:
            self._group_id = uuid.uuid4()
            self._path = (self._name,)
        self._cv_token = _current_logger.set(self)
        if self._enter_msg:
            if not isinstance(self._enter_msg, str):
                self.log(
                    self._ctx_level, "entered: %s", self._name, extra={EXTRA_KEY: {ENTER_KEY: True}}
                )
            else:
                self.log(self._ctx_level, self._enter_msg, extra={EXTRA_KEY: {ENTER_KEY: True}})
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Exit the context.

        The logger will only be removed from the hierarchy if it is exiting
        for the last time, i.e., the number of times it has been entered,
        and exited are balanced.

        If the logger is removed from the hierarchy, or it is exiting
        due to an exception, an exited message will be logged. See
        :meth:`~Logger.__enter__` for how the message is logged, and how to
        suppress it.  The message text will be ``"exited: %s" % name``.
        """
        self._entered -= 1
        exc_info = (exc_type, exc_val, exc_tb) if exc_type and exc_val else None
        if (self._entered == 0 or exc_info) and self._exit_msg:
            if not isinstance(self._exit_msg, str):
                self.log(
                    self._ctx_level,
                    "exited: %s",
                    self._name,
                    exc_info=exc_info,
                    extra={EXTRA_KEY: {EXIT_KEY: True}},
                )
            else:
                self.log(
                    self._ctx_level,
                    self._exit_msg,
                    exc_info=exc_info,
                    extra={EXTRA_KEY: {EXIT_KEY: True}},
                )
        if self._entered == 0 and self._cv_token:
            _current_logger.reset(self._cv_token)

    @property
    def collectors(self) -> Mapping[str, Collector[Any]]:
        """The logger collectors."""
        return self._collectors

    @property
    def group_id(self) -> uuid.UUID | None:
        """The group ID for the hierarchy assigned by the root logger."""
        return self._group_id

    @property
    def name(self) -> str:
        """The logger name.

        This is not the name of the underlying Python logger. This is the name
        of the logger in the path if the logger has been entered. Otherwise,
        it is the ``name`` that the logger was initialized with.
        """
        return self._path[-1] if self._path else self._name

    @property
    def parent(self) -> Logger | None:
        """The parent logger, or None if this is a root logger."""
        return self._parent

    @property
    def path(self) -> tuple[str, ...] | None:
        """The logger path.

        The path is a sequence of logger names following the hierarchy of
        loggers from the root logger to this one. The final name in the
        path will be the name that the logger was initialized with if a
        logger with the same name has not been registered at that path index
        yet. Otherwise, an incremental number is appended to the initial
        name to get a unique name for that path index.

        Example:
            If there are 2 loggers named "baz" at the third level of the
            hierarchy, their respecitve paths will be:

            .. code-block:: python

                ["foo", "bar", "baz"]
                ["foo", "bar", "baz2"]
        """
        return self._path

    def register_path(self, logger: Logger) -> tuple[str, ...]:
        """A logger can register a path for itself with its parent.

        This method should not be called directly. The logger will
        automatically register itself with its parent.
        """
        if logger.parent is not self:
            msg = "`logger` is not a child of this logger"
            raise RuntimeError(msg)
        if not self._path:  # pragma: no cover
            msg = "path has not been initialized"
            raise RuntimeError(msg)
        name = logger.name
        if name not in self._path_registry:
            self._path_registry[name] = 1
            return (*self._path, name)
        self._path_registry[name] += 1
        return (*self._path, f"{name}{self._path_registry[name]}")

    def collect(
        self,
        *identifiers: str | uuid.UUID,
        log_item: LogItem,
        extra_tremors: MutableMapping[str, object] | None = None,
    ) -> MutableMapping[str, object]:
        """Run the specified collectors.

        Args:
            *identifiers: Collector names, or IDs to run.
            log_item: The LogItem to pass to collectors.
            extra_tremors: The initial ``EXTRA_KEY`` extras. This mapping is mutated, and the new
                collector states are added to it. If none, an empty dict will be used.

        Returns:
            The mutated ``extra_tremors``, or a new dict with the collector states.
        """
        if not identifiers:
            collectors_to_collect: Iterable[Collector[Any]] = self._collectors.values()
        else:
            inital_names: set[str] = set()
            inital_ids: set[uuid.UUID] = set()
            name_set, id_set = functools.reduce(
                _names_and_ids_reducer, identifiers, (inital_names, inital_ids)
            )
            collectors_to_collect = filter(
                lambda c: c.name in name_set or c.id in id_set, self._collectors.values()
            )
        initial_collectors: MutableMapping[str, Collector[Any]] = {}
        if extra_tremors is None:
            extra_tremors = {}
        self._collectors, extra_tremors = functools.reduce(
            _collectors_reducer,
            zip(itertools.repeat(log_item), collectors_to_collect, strict=False),
            (initial_collectors, extra_tremors),
        )
        return extra_tremors

    def log(  # noqa: PLR0913
        self,
        level: int,
        msg: object,
        *args: object,
        exc_info: bool
        | tuple[type[BaseException], BaseException, TracebackType | None]
        | tuple[None, None, None]
        | BaseException
        | None = None,
        stack_info: bool = False,
        stacklevel: int = 1,
        extra: Mapping[str, object] | None = None,
        **_kwargs: object,  # needed for the type checker
    ) -> None:
        """Run collectors, and log a message at the given level.

        The arguments are interpreted as for :meth:`~logging.Logger.debug`.

        This method mutates, or replaces ``extra[EXTRA_KEY]``.
        """
        some_extra = extra or {}
        extra_tremors = some_extra.get(EXTRA_KEY)
        if not isinstance(extra_tremors, MutableMapping):
            extra_tremors = {}
        log_item = LogItem(
            logger=self,
            level=level,
            msg=msg,
            args=args,
            exc_info=exc_info,
            stack_info=stack_info,
            stacklevel=stacklevel,
            extra=extra,
            record=None,
        )
        self.collect(log_item=log_item, extra_tremors=extra_tremors)
        return self.logger.log(
            level,
            msg,
            *args,
            exc_info=exc_info,
            stack_info=stack_info,
            stacklevel=stacklevel,
            extra={**some_extra, EXTRA_KEY: extra_tremors},
        )


def enrich_std_record(record: logging.LogRecord) -> None:
    """Enrich a record not created by a tremors logger.

    This function mutates ``record`` if the record doesn't have the tremors
    attribute, ``EXTRA_KEY``, i.e., a standard logger produced the record. The
    current logger's collectors are run on a LogItem derived from the record's
    attributes. Then the results are stored on the record in the ``EXTRA_KEY``
    attribute, so the record may subsequently be processed by collector
    formatters.

    The function may be used in filters to capture, and treat standard logs, as
    if they were logged by the current logger.

    If there is no current logger, the record won't be modified.
    """
    if hasattr(record, EXTRA_KEY) or not (logger := _current_logger.get()):
        return

    log_item = LogItem(
        logger=logger,
        level=record.levelno,
        msg=getattr(record, "message", None) or record.getMessage(),
        args=record.args,
        exc_info=record.exc_info,
        stack_info=bool(getattr(record, "stack_info", None)),
        stacklevel=None,
        extra=None,
        record=record,
    )
    record.__dict__[EXTRA_KEY] = logger.collect(log_item=log_item, extra_tremors={})
