"""This module contains collector bundles and tools.

A collector bundle is an object that contains the definition of
the collector state, a factory function to create a collector, and a
formatter function that will extract, and format the collector state from
a :class:`~logging.LogRecord`.

In addition to the bundles in this module, you can create custom bundles. A
collector is generic in the type of its state, which can be of an arbitrary
type. A collector factory can be defined to have arbitrary parameters that
depend on how the collector is initialized. A collector formatter has a
fixed signature as described in :class:`Formatter`.
"""

import dataclasses
import logging
import sys
import time
from collections.abc import Mapping
from typing import TYPE_CHECKING, NamedTuple, Protocol

import tremors.logger

if TYPE_CHECKING:
    import uuid
    from collections.abc import Callable
elif "sphinx" in sys.modules:
    import uuid  # noqa: TC003
    from collections.abc import Callable


def get_state[T](typ: type[T], name: str, record: logging.LogRecord) -> T | None:
    """Get the collector state from a :class:`~logging.LogRecord`.

    The function can be used by custom collector formatters to get the
    collector state from a record.

    Args:
        typ: The collector state class.
        name: The name of the collector.
        record: The LogRecord the may contain the collector state.

    Returns:
        An object of type ``typ`` if the LogRecord contains such an object
        for the collector. Otherwise None.
    """
    tremors_extra = getattr(record, tremors.logger.EXTRA_KEY, None)
    if isinstance(tremors_extra, Mapping):
        state = tremors_extra.get(name)
        if isinstance(state, typ):
            return state
    return None


class Formatter(Protocol):
    """Callback proctocol for a collector formatter."""

    def __call__(
        self, record: logging.LogRecord, *, name: str = "", fmt: str = "", default: str = ""
    ) -> str:
        """Extract and format a collector state.

        Args:
            record: A :class:`~logging.LogRecord` that may contain the
                collector state.
            name: The name of the collector.
            fmt: A format specification. Implementations commonly expect
                this to adhere to the :ref:`formatspec`. Then ``str(fmt)``
                may be used to obtain the specification, so ``fmt`` does
                not have to be a string.
            default: A value that may be returned if the formatter could
                not obtain, or format the state.

        Returns:
            The state formatted as a string, or the default value.
        """


class CollectorBundle[T, **P](NamedTuple):
    """A bundle to create, and work with a collector."""

    state: type[T]
    factory: Callable[P, tremors.logger.Collector[T]]
    formatter: Formatter


class IndentifierState(NamedTuple):
    """The state of a collector created by :func:`identifier_factory`."""

    group_id: uuid.UUID
    parent: tremors.logger.Logger | None
    path: tuple[str, ...]


def _identifier_collect(
    state: IndentifierState | None, log_item: tremors.logger.LogItem
) -> IndentifierState | None:
    if state:
        return state
    logger = log_item.logger
    if not logger.group_id or not logger.path:
        return None
    return IndentifierState(group_id=logger.group_id, parent=logger.parent, path=logger.path)


def identifier_factory(
    name: str = "identifier", *, level: int = logging.INFO
) -> tremors.logger.Collector[IndentifierState | None]:
    """Create a collector to gather logger identifiers.

    Args:
        name: The collector name.
        level: The collector level.

    Returns:
        An identifier collector with the specified name, and level.
    """
    return tremors.logger.Collector(name=name, level=level, state=None, collect=_identifier_collect)


def identifier_formatter(
    record: logging.LogRecord,
    *,
    name: str = "identifier",
    fmt: str = "[{process}{thread}{group_id}] {path}",
    default: str = "",
) -> str:
    """Extract an :class:`IndentifierState` from a record, and format it.

    ``str(fmt)`` may contain any of the fields in this table. The formatted
    string is generated by replacing the fields with the values described
    in the table.

    ========  =====
    Field     Value
    ========  =====
    process   ``f"p:{pid}"`` where ``pid`` is the process ID if it's
               available (notice the trailing space). Otherwise the empty
               string.
    thread    ``f"t:{tid}"`` where ``tid`` is the thread ID if it's
               available (notice the trailing space). Otherwise the empty
               string.
    group_id  ``f"t:{gid}"`` where ``gid`` is the group ID.
    path      The path represented as a POSIX-like path.
    ========  =====

    Returns:
        The formatted string if the record contains a state for the
        collector. Otherwise ``default``.
    """
    state = get_state(IndentifierState, name, record)

    if state is None:
        return default

    process = f"p:{record.process} " if record.process is not None else ""
    thread = f"t:{record.thread} " if record.thread is not None else ""
    group_id = f"g:{state.group_id}"
    path = "/".join(state.path)
    return str(fmt).format(process=process, thread=thread, group_id=group_id, path=path)


identifier = CollectorBundle(
    state=IndentifierState, factory=identifier_factory, formatter=identifier_formatter
)
"""The identifier collector bundle.

Examples:

    With :func:`~logging.basicConfig` we configure the standard root logger
    with a stream handler that has a formatter that incorporates the record's
    ``ident`` attribute in the message. We add the filter function ``flt`` to
    the handler. The filter uses the formatter from the identifier collector
    bundle to extract the collector state from a record, and format it.
    Then the filter returns a modified record with the ident attribute set
    to the formatted state.

    .. note::

        If we had added the filter to the logger instead of the handler,
        the filter would not be called for messages logged by descendant
        standard loggers.  But by adding the filter to the handler, we ensure
        that it is called for all messages that the handler could emit.

    .. code-block:: python

        import copy
        import logging

        import tremors
        from tremors import collector


        def flt(record: logging.LogRecord) -> logging.LogRecord:
            record = copy.copy(record)
            ident = collector.identifier.formatter(record)
            record.ident = f"{ident} > " if ident else ""
            return record


        @tremors.logged(collector.identifier.factory())
        def parent(*, logger: tremors.Logger = tremors.from_logged) -> None:
            logger.info("foo")
            child()


        @tremors.logged(collector.identifier.factory())
        def child(*, logger: tremors.Logger = tremors.from_logged) -> None:
            logger.info("bar")


        logging.basicConfig(
            format="%(ident)s%(levelname)s:%(name)s:%(message)s", level=logging.INFO
        )
        logging.root.handlers[0].addFilter(flt)
        parent()

    The logged messages contain the process ID, thread ID, group ID, and
    hierarchy as a path.

    .. code-block:: shell

        [p:... t:... g:...] parent > INFO:root:entered: parent
        [p:... t:... g:...] parent > INFO:root:foo
        [p:... t:... g:...] parent/child > INFO:root:entered: child
        [p:... t:... g:...] parent/child > INFO:root:bar
        [p:... t:... g:...] parent/child > INFO:root:exited: child
        [p:... t:... g:...] parent > INFO:root:exited: parent

    We can expand on the previous example by throwing a function that uses
    standard logging into the mix. We modify ``child`` to call ``other``,
    which could be a function that you can't change, and doesn't use Tremors
    loggers. By default, standard loggers propagate their messages, so
    descendant messages will eventually be processed by the root logger's
    handler with our filter. The collector formatter will not be able to
    extract the state from a descendant record, and will return the default
    value, the empty string. Then ``record.ident`` will be the empty string,
    which will be interpolated into the formatted message.

    .. note::

        Since the formatter expects the record to have an ident attribute,
        the filter must always add this attribute to the record.

    .. code-block:: python

        @tremors.logged(collector.identifier.factory())
        def child(*, logger: tremors.Logger = tremors.from_logged) -> None:
            logger.info("bar")
            other()


        def other() -> None:
            logging.getLogger(__name__).info("baz")


        parent()

    The functions that use Tremors loggers have entering and exiting
    messages, and their messages contain identifier information, as before.
    The ``other`` function, which logs with standard loggers simply emits
    messages with empty ident sections.

    .. code-block:: shell

        [p:... t:... g:...] parent > INFO:root:entered: parent
        [p:... t:... g:...] parent > INFO:root:foo
        [p:... t:... g:...] parent/child > INFO:root:entered: child
        [p:... t:... g:...] parent/child > INFO:root:bar
        INFO:__main__:baz
        [p:... t:... g:...] parent/child > INFO:root:exited: child
        [p:... t:... g:...] parent > INFO:root:exited: parent
"""


@dataclasses.dataclass
class CounterState:
    """The state of a collector created by :func:`identifier_factory`.

    This state is mutable so that a single collector instance can be shared
    with multiple loggers, and each logger can update the same state.
    """

    count: int
    step: int


def _counter_collect(state: CounterState, _: tremors.logger.LogItem) -> CounterState:
    state.count += state.step
    return state


def counter_factory(
    name: str = "counter", *, level: int = logging.INFO, initial: int = 0, step: int = 1
) -> tremors.logger.Collector[CounterState]:
    """Create a collector to count logging calls.

    Args:
        name: The collector name.
        level: The collector level.
        initial: The initial value of the count.
        step: How much to increase the count by each time the collector runs.

    Returns:
        A counter collector with the specified name, level, behavior, and mutable state.
    """
    return tremors.logger.Collector(
        name=name,
        level=level,
        state=CounterState(count=initial, step=step),
        collect=_counter_collect,
    )


def counter_formatter(
    record: logging.LogRecord, *, name: str = "counter", fmt: str = "{counter}", default: str = "0"
) -> str:
    """Extract a count from a record, and format it.

    ``str(fmt)`` may contain the field ``counter``. The formatted string is
    generated by replacing the field with the state's count.

    Returns:
        The formatted string if the record contains a state for the
        collector. Otherwise ``default``.
    """
    state = get_state(CounterState, name, record)

    if state is None:
        return default

    return str(fmt).format(counter=state.count)


counter = CollectorBundle(state=CounterState, factory=counter_factory, formatter=counter_formatter)
"""The counter collector bundle.

Examples:

    We configure standard logging in a similar manner to the examples in
    :attr:`identifier`, except here we use counter collectors.

    .. code-block:: python

        import copy
        import logging

        import tremors
        from tremors import collector


        def flt(record: logging.LogRecord) -> logging.LogRecord:
            record = copy.copy(record)
            count = collector.counter.formatter(record, fmt="[{counter}] ", default="")
            record.count = count
            return record


        @tremors.logged(collector.counter.factory())
        def parent(*, logger: tremors.Logger = tremors.from_logged) -> None:
            logger.info("foo")
            child()


        @tremors.logged(collector.counter.factory())
        def child(*, logger: tremors.Logger = tremors.from_logged) -> None:
            logger.info("bar")


        logging.basicConfig(
            format="%(count)s%(levelname)s:%(name)s:%(message)s", level=logging.INFO
        )
        logging.root.handlers[0].addFilter(flt)
        parent()

    The logged messages include the count, where the ``parent`` and ``child``
    counts are distinct from each other because each loggers uses its own
    counter collector.

    .. code-block:: shell

        [1] INFO:root:entered: parent
        [2] INFO:root:foo
        [1] INFO:root:entered: child
        [2] INFO:root:bar
        [3] INFO:root:exited: child
        [3] INFO:root:exited: parent

    Now let's modify the previous example to have ``parent`` and ``child``
    use a shared counter.

    .. code-block:: python

        shared_counter = collector.counter.factory()


        @tremors.logged(shared_counter)
        def parent(*, logger: tremors.Logger = tremors.from_logged) -> None:
            logger.info("foo")
            child()


        @tremors.logged(shared_counter)
        def child(*, logger: tremors.Logger = tremors.from_logged) -> None:
            logger.info("bar")


        parent()

    Here the count is shared between the two functions.

    .. code-block:: shell

        [1] INFO:root:entered: parent
        [2] INFO:root:foo
        [3] INFO:root:entered: child
        [4] INFO:root:bar
        [5] INFO:root:exited: child
        [6] INFO:root:exited: parent
"""


@dataclasses.dataclass
class ElapsedState:
    """The state of an :func:`elapsed` collector.

    The state is mutable like :class:`CounterState`.
    """

    t0: int | None
    t: int | None


def _collect_elapsed(state: ElapsedState, _: tremors.logger.LogItem) -> ElapsedState:
    t = time.perf_counter_ns()
    if state.t0 is None:
        state.t0 = t
    state.t = t
    return state


def elapsed_factory(
    name: str = "elapsed", *, level: int = logging.INFO
) -> tremors.logger.Collector[ElapsedState]:
    """Create a collector to measure how much time has elapsed.

    Args:
        name: The collector name.
        level: The collector level.

    Returns:
        An elapsed collector with the specified name, and level.
    """
    return tremors.logger.Collector(
        name=name, level=level, state=ElapsedState(t0=None, t=None), collect=_collect_elapsed
    )


def elapsed_formatter(
    record: logging.LogRecord,
    *,
    name: str = "elapsed",
    fmt: str = "{elapsed:.3f}",
    default: str = "",
) -> str:
    """Extract an :class:`ElapsedState` from a record, and format it.

    ``str(fmt)`` may contain the field ``elapsed``. The formatted string is
    generated by replacing the field with the state's elapsed time in seconds.

    Returns:
        The formatted string if the record contains a state for the
        collector. Otherwise ``default``.
    """
    state = get_state(ElapsedState, name, record)

    if state is None:
        return default

    t0, t = state.t0, state.t

    if t0 is None or t is None:
        return default

    elapsed_s = (t - t0) / 1e9
    return str(fmt).format(elapsed=elapsed_s)


elapsed = CollectorBundle(state=ElapsedState, factory=elapsed_factory, formatter=elapsed_formatter)
"""The elapsed collector bundle.

Examples:

    We configure standard logging in a similar manner to the examples in
    :attr:`identifier`, except here we use elapsed collectors, and identifier
    collectors with a custom format to only show the path.

    .. code-block:: python

        import copy
        import logging
        import time

        import tremors
        from tremors import collector


        def flt(record: logging.LogRecord) -> logging.LogRecord:
            record = copy.copy(record)
            ident = collector.identifier.formatter(record, fmt="{path}")
            elapsed = collector.elapsed.formatter(record)
            record.ident = f"{ident} > " if ident else ""
            record.elapsed = f"{elapsed} " if elapsed else ""
            return record


        @tremors.logged(collector.identifier.factory(), collector.elapsed.factory())
        def parent(*, logger: tremors.Logger = tremors.from_logged) -> None:
            logger.info("sleeping for 1s...")
            time.sleep(1)
            child()


        @tremors.logged(collector.identifier.factory(), collector.elapsed.factory())
        def child(*, logger: tremors.Logger = tremors.from_logged) -> None:
            logger.info("sleeping for 1s...")
            time.sleep(1)


        logging.basicConfig(
            format="%(elapsed)s%(ident)s%(levelname)s:%(name)s:%(message)s",
            level=logging.INFO,
        )
        logging.root.handlers[0].addFilter(flt)
        parent()

    The logged messages include the elapsed time since each elapsed collector
    started.

    .. code-block:: shell

        0.000 parent > INFO:root:entered: parent
        0.000 parent > INFO:root:sleeping for 1s...
        0.000 parent/child > INFO:root:entered: child
        0.000 parent/child > INFO:root:sleeping for 1s...
        1.001 parent/child > INFO:root:exited: child
        2.001 parent > INFO:root:exited: parent

    We can try the same thing again, only this time with a shared elapsed
    collector.

    .. code-block:: python

        shared_elapsed = collector.elapsed.factory()


        @tremors.logged(collector.identifier.factory(), shared_elapsed)
        def parent(*, logger: tremors.Logger = tremors.from_logged) -> None:
            logger.info("sleeping for 1s...")
            time.sleep(1)
            child()


        @tremors.logged(collector.identifier.factory(), shared_elapsed)
        def child(*, logger: tremors.Logger = tremors.from_logged) -> None:
            logger.info("sleeping for 1s...")
            time.sleep(1)


        parent()

    Notice that this time, the ``child`` elapsed time starts at 1s instead
    of resetting to 0s.

    .. code-block:: shell

        0.000 parent > INFO:root:entered: parent
        0.000 parent > INFO:root:sleeping for 1s...
        1.000 parent/child > INFO:root:entered: child
        1.001 parent/child > INFO:root:sleeping for 1s...
        2.001 parent/child > INFO:root:exited: child
        2.001 parent > INFO:root:exited: parent
"""
