"""Logging utilities for pretty-printing, colorized, and context-managed logging.

This module provides:
- ANSI color and style helpers for log output
- A custom Formatter for enhanced log formatting
- Functions to initialize and configure logging
- Context manager for temporary log level changes
"""

import contextlib
import logging
import os
import re
from datetime import UTC, datetime


def colorize(text, color_code):
    """Wraps text with the ANSI escape code for the given color."""
    return f"\033[{color_code}m{text}\033[39m"


def bold(text):
    """Bolds the given text."""
    return f"\033[1m{text}\033[22m"


def faint(text):
    """Faints the given text."""
    return f"\033[2m{text}\033[22m"


def italic(text):
    """Italicizes the given text."""
    return f"\033[3m{text}\033[23m"


def black(text):
    """Colors the given text black using ANSI escape codes."""
    return colorize(text, 30)


def red(text):
    """Colors the given text red using ANSI escape codes."""
    return colorize(text, 31)


def green(text):
    """Colors the given text green using ANSI escape codes."""
    return colorize(text, 32)


def yellow(text):
    """Colors the given text yellow using ANSI escape codes."""
    return colorize(text, 33)


def blue(text):
    """Colors the given text blue using ANSI escape codes."""
    return colorize(text, 34)


def magenta(text):
    """Colors the given text magenta using ANSI escape codes."""
    return colorize(text, 35)


def cyan(text):
    """Colors the given text cyan using ANSI escape codes."""
    return colorize(text, 36)


def white(text):
    """Colors the given text white using ANSI escape codes."""
    return colorize(text, 37)


CRITICAL = logging.CRITICAL
DEBUG = logging.DEBUG
ERROR = logging.ERROR
INFO = logging.INFO
TRACE = 5
WARNING = logging.WARNING

LEVEL_LABELS = {
    "CRITICAL": red(bold("CRITICAL")),
    "DEBUG": magenta("DEBUG"),
    "ERROR": red("ERROR"),
    "INFO": white("INFO"),
    "TRACE": green(faint(bold("TRACE"))),
    "WARNING": yellow("WARNING"),
}


def strip_ansi(s: str) -> str:
    """Removes ANSI escape sequences from the given string."""
    r = re.compile(r"\x1B\[.*?[a-zA-Z]")
    return r.sub("", s)


def disp_len(s: str) -> int:
    """Returns the display length of the given string."""
    return len(strip_ansi(s))


def indent_rest(input_string: str, indent: int) -> str:
    """Indents all but the first line of the given string."""
    lines = input_string.split("\n")
    return "\n".join([lines[0]] + [f"{' ' * indent}{line}" for line in lines[1:]])


class Formatter(logging.Formatter):
    """Pretty-printing log formatter."""

    def format(self, record):
        """Format the specified record as a pretty-printed, colorized log message.

        Parameters
        ----------
        record : logging.LogRecord
            The log record to be formatted.

        Returns:
        -------
        str
            The formatted log message string.
        """
        ts = (
            datetime.fromtimestamp(record.created, UTC).strftime(
                "%Y-%m-%dT%H:%M:%S.%f"
            )[:-3]
            + "Z"
        )
        level = LEVEL_LABELS.get(record.levelname, record.levelname)
        n = record.name
        color = sum([ord(x) for x in n]) % 6 + 32
        msg = record.getMessage()
        base = f"{italic(ts)} - {colorize(n, color)} - {level} - "
        ex = "\n" + self.formatException(record.exc_info) if record.exc_info else ""
        w = disp_len(base)
        return f"{base}{indent_rest(msg, w)}{indent_rest(ex, w)}"

    def format_exception(self, exc_info):
        """Formats and returns the exception information as a string.

        Parameters
        ----------
        exc_info : tuple
            Exception information as returned by sys.exc_info().

        Returns:
        -------
        str
            The formatted exception string.
        """
        return super().formatException(exc_info)


def get_logger(name):
    """Returns a logger with the specified name.

    Parameters
    ----------
    name : str
        The name of the logger to retrieve.

    Returns:
    -------
    logging.Logger
        The logger instance associated with the given name.
    """
    return logging.getLogger(name)


def set_level(level):
    """Sets the global log level."""
    # TODO: support setting trace
    if (isinstance(level, int) and level > 0) or level in logging._nameToLevel:
        logging.getLogger().setLevel(level)
    else:
        logging.getLogger(__name__).warning(
            "Invalid log level %s; ignoring.", red(level)
        )


def init(level: str | int | None = None):
    """Initializes the logging system."""
    logging.captureWarnings(True)
    level = level or os.environ.get("LOG_LEVEL", "INFO").upper()
    logger = logging.getLogger()
    handler = logging.StreamHandler()
    formatter = Formatter("%(message)s")
    handler.setFormatter(formatter)
    logger.handlers = [handler]
    set_level(level)


@contextlib.contextmanager
def level(level):
    """Temporarily modifies the log level.

    Usage:
    ```
    with log.level(log.DEBUG):
        ...
    ```.
    """
    logger = logging.getLogger()
    old_level = logger.level
    logger.setLevel(level)
    try:
        yield  # Yield control back to the calling block
    finally:
        logger.setLevel(old_level)  # Restore the original log level
