import sys
import logging
import time
import colorlog

# A module is only ever loaded once, so we can set up the logger here
logger = logging.getLogger("custom")

# Set the default logging level once at the module level
TRACE = 5
logger.setLevel(TRACE)


def trace(self, message, *args, **kwargs):
    """Log a message with severity TRACE."""
    if self.isEnabledFor(TRACE):
        self._log(TRACE, message, args, **kwargs)


def setup_logger():
    """Set up the logger with a custom format and color scheme."""
    # Set the log record to include padding for the origin
    old_factory = logging.getLogRecordFactory()

    def record_factory(*args, **kwargs):
        record = old_factory(*args, **kwargs)
        record.origin = f'{record.filename}:{record.lineno}'
        return record
    logging.setLogRecordFactory(record_factory)

    # Create a new log level
    logging.Logger.trace = trace
    logging.addLevelName(TRACE, "TRACE")

    # Set formatting for the custom logger
    fmt = colorlog.ColoredFormatter(
        "%(asctime)s | %(log_color)s%(levelname)-5s%(reset)s | %(blue)s%(origin)-20s%(reset)s | %(log_color)s%(message)s%(reset)s",
        datefmt="%H:%M:%S",
        log_colors={
            'TRACE': 'yellow',
            'DEBUG': 'white',
            'INFO': 'green',
            'WARNING': 'yellow',
            'ERROR': 'red',
        }
    )

    stdout = colorlog.StreamHandler(stream=sys.stdout)
    stdout.setFormatter(fmt)
    logger.addHandler(stdout)

    # Disable the default logging handler
    logger.propagate = False


def log_time(func):
    """Decorator to log the time taken by a function."""
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        logger.trace(
            f"{func.__name__} completed in {time.time() - start:.2f} seconds")
        return result
    return wrapper


# Actually set up the logger for all modules that import this one
setup_logger()