"""This module contains decorators for using Tremors loggers.

The :func:`logged` decorator wraps a callable in a Tremors logger context,
and injects that logger into the callable.
"""

import functools
import logging
from collections.abc import Callable, Coroutine
from typing import Any, overload

import tremors.logger


def _logged[TRet, **P](  # noqa: PLR0913
    fn: Callable[P, TRet],
    *collectors: tremors.logger.Collector[Any] | tremors.logger.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,
) -> Callable[P, TRet]:
    @functools.wraps(fn)
    def logged_wrapper(*args: P.args, **kwargs: P.kwargs) -> TRet:
        if (logger_arg := kwargs.get("logger")) and logger_arg is not from_logged:
            return fn(*args, **kwargs)
        with tremors.logger.Logger(
            *collectors,
            name=name,
            logger_name=logger_name,
            ctx_level=ctx_level,
            enter_msg=enter_msg,
            exit_msg=exit_msg,
            is_root=is_root,
        ) as logger:
            kwargs["logger"] = logger
            return fn(*args, **kwargs)

    return logged_wrapper


@overload
def logged[TRet, **P](fn_or_collector: Callable[P, TRet]) -> Callable[P, TRet]: ...


@overload
def logged[TRet, **P](
    fn_or_collector: tremors.logger.Collector[Any] | tremors.logger.CollectorFactory[Any],
    *collectors: tremors.logger.Collector[Any] | tremors.logger.CollectorFactory[Any],
    name: str | None = None,
    logger_name: str | None = None,
    ctx_level: int = logging.INFO,
    enter_msg: str | bool = True,
    exit_msg: str | bool = True,
    is_root: bool = False,
) -> Callable[[Callable[P, TRet]], Callable[P, TRet]]: ...


@overload
def logged[TRet, **P](
    *,
    name: str | None = None,
    logger_name: str | None = None,
    ctx_level: int = logging.INFO,
    enter_msg: str | bool = True,
    exit_msg: str | bool = True,
    is_root: bool = False,
) -> Callable[[Callable[P, TRet]], Callable[P, TRet]]: ...


def logged[TRet, **P](  # noqa: PLR0913
    fn_or_collector: Callable[P, TRet]
    | tremors.logger.Collector[Any]
    | tremors.logger.CollectorFactory[Any]
    | None = None,
    *collectors: tremors.logger.Collector[Any] | tremors.logger.CollectorFactory[Any],
    name: str | None = None,
    logger_name: str | None = None,
    ctx_level: int = logging.INFO,
    enter_msg: str | bool = True,
    exit_msg: str | bool = True,
    is_root: bool = False,
) -> Callable[P, TRet] | Callable[[Callable[P, TRet]], Callable[P, TRet]]:
    """Inject a :class:`~tremors.logger.Logger` into a callable.

    The callable must have a keyword-only parameter ``logger``. If the
    value for ``logger`` is the default value, :data:`from_logged`, it is
    replaced by a logger with the specified collectors, name, and underlying
    logger. This logger is entered, and exited immediately before and after
    the callable runs. Otherwise, the supplied logger will be used in the
    call. A supplied logger will not be automatically entered, and exited.

    If ``fn_or_collector`` is a callable, e.g., the decorator was used like
    ``@logged``, it is assumed that no collectors were specified, and the
    default collectors are used. Otherwise, it is assumed this is the first
    collector, e.g., the decorator was used like ``@logged(some_collector,
    ...)``.

    ``*collectors`` is the rest of the collectors, with the first collector
    being ``fn_or_collector``.

    The rest of the arguments are passed to :class:~`tremors.logger.Logger`.
    """
    if callable(fn_or_collector) and not isinstance(
        fn_or_collector, tremors.logger.CollectorFactory
    ):
        default_collectors = ()  # TODO @NAS: Implement default collectors.
        return _logged(fn_or_collector, *default_collectors, name=fn_or_collector.__name__)

    def decorator(fn: Callable[P, TRet]) -> Callable[P, TRet]:
        return (
            _logged(
                fn,
                fn_or_collector,
                *collectors,
                name=name if name is not None else fn.__name__,
                logger_name=logger_name,
                ctx_level=ctx_level,
                enter_msg=enter_msg,
                exit_msg=exit_msg,
                is_root=is_root,
            )
            if fn_or_collector
            else _logged(
                fn,
                *collectors,
                name=name if name is not None else fn.__name__,
                logger_name=logger_name,
                ctx_level=ctx_level,
                enter_msg=enter_msg,
                exit_msg=exit_msg,
                is_root=is_root,
            )
        )

    return decorator


def _async_logged[TRet, **P](  # noqa: PLR0913
    fn: Callable[P, Coroutine[object, object, TRet]],
    *collectors: tremors.logger.Collector[Any] | tremors.logger.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,
) -> Callable[P, Coroutine[object, object, TRet]]:
    @functools.wraps(fn)
    async def logged_wrapper(*args: P.args, **kwargs: P.kwargs) -> TRet:
        if (logger_arg := kwargs.get("logger")) and logger_arg is not from_logged:
            return await fn(*args, **kwargs)
        with tremors.logger.Logger(
            *collectors,
            name=name,
            logger_name=logger_name,
            ctx_level=ctx_level,
            enter_msg=enter_msg,
            exit_msg=exit_msg,
            is_root=is_root,
        ) as logger:
            kwargs["logger"] = logger
            return await fn(*args, **kwargs)

    return logged_wrapper


@overload
def async_logged[TRet, **P](
    fn_or_collector: Callable[P, Coroutine[object, object, TRet]],
) -> Callable[P, Coroutine[object, object, TRet]]: ...


@overload
def async_logged[TRet, **P](
    fn_or_collector: tremors.logger.Collector[Any] | tremors.logger.CollectorFactory[Any],
    *collectors: tremors.logger.Collector[Any] | tremors.logger.CollectorFactory[Any],
    name: str | None = None,
    logger_name: str | None = None,
    ctx_level: int = logging.INFO,
    enter_msg: str | bool = True,
    exit_msg: str | bool = True,
    is_root: bool = False,
) -> Callable[
    [Callable[P, Coroutine[object, object, TRet]]], Callable[P, Coroutine[object, object, TRet]]
]: ...


@overload
def async_logged[TRet, **P](
    *,
    name: str | None = None,
    logger_name: str | None = None,
    ctx_level: int = logging.INFO,
    enter_msg: str | bool = True,
    exit_msg: str | bool = True,
    is_root: bool = False,
) -> Callable[
    [Callable[P, Coroutine[object, object, TRet]]], Callable[P, Coroutine[object, object, TRet]]
]: ...


def async_logged[TRet, **P](  # noqa: PLR0913
    fn_or_collector: Callable[P, Coroutine[object, object, TRet]]
    | tremors.logger.Collector[Any]
    | tremors.logger.CollectorFactory[Any]
    | None = None,
    *collectors: tremors.logger.Collector[Any] | tremors.logger.CollectorFactory[Any],
    name: str | None = None,
    logger_name: str | None = None,
    ctx_level: int = logging.INFO,
    enter_msg: str | bool = True,
    exit_msg: str | bool = True,
    is_root: bool = False,
) -> (
    Callable[P, Coroutine[object, object, TRet]]
    | Callable[
        [Callable[P, Coroutine[object, object, TRet]]], Callable[P, Coroutine[object, object, TRet]]
    ]
):
    """Inject a :class:`~tremors.logger.Logger` into a async callable.

    Apart for the async callable, the arguments are the same as for
    :func:`logged`.
    """
    if callable(fn_or_collector) and not isinstance(
        fn_or_collector, tremors.logger.CollectorFactory
    ):
        default_collectors = ()  # TODO @NAS: Implement default collectors.
        return _async_logged(fn_or_collector, *default_collectors, name=fn_or_collector.__name__)

    def decorator(
        fn: Callable[P, Coroutine[object, object, TRet]],
    ) -> Callable[P, Coroutine[object, object, TRet]]:
        return (
            _async_logged(
                fn,
                fn_or_collector,
                *collectors,
                name=name if name is not None else fn.__name__,
                logger_name=logger_name,
                ctx_level=ctx_level,
                enter_msg=enter_msg,
                exit_msg=exit_msg,
                is_root=is_root,
            )
            if fn_or_collector
            else _async_logged(
                fn,
                *collectors,
                name=name if name is not None else fn.__name__,
                logger_name=logger_name,
                ctx_level=ctx_level,
                enter_msg=enter_msg,
                exit_msg=exit_msg,
                is_root=is_root,
            )
        )

    return decorator


from_logged = tremors.logger.Logger(name="__from_logged__")
"""A sentinel logger value.

When used as the ``logger`` argument to a :func:`logged` callable, the
callable is wrapped in a :class:`~tremors.logger.Logger` context that is
injected into the callable as the ``logger`` argument.
"""
