import logging
from logging.handlers import TimedRotatingFileHandler
from typing import Any, Dict, List

from rich.logging import RichHandler
from rich.traceback import install as rich_tb_install

from coalescenceml.constants import APP_NAME, COML_LOGGING_VERBOSITY
from coalescenceml.enums import LoggingLevels


LOG_FILE = f"{APP_NAME}_logs.log"


class CustomFormatter(logging.Formatter):
    """Formats logs according to our specification."""

    grey: str = "\x1b[38;21m"
    pink: str = "\x1b[35m"
    green: str = "\x1b[32m"
    blue: str = "\x1b[34m"
    yellow: str = "\x1b[33m"
    red: str = "\x1b[31m"
    bold_red: str = "\x1b[31;1m"
    purple: str = "\x1b[1;35m"
    reset: str = "\x1b[0m"

    format_template: str = (
        "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%("
        "filename)s:%(lineno)d)"
        if LoggingLevels[COML_LOGGING_VERBOSITY] == LoggingLevels.DEBUG
        else "%(message)s"
    )

    COLORS: Dict[LoggingLevels, str] = {
        LoggingLevels.DEBUG: purple,
        LoggingLevels.INFO: blue,
        LoggingLevels.WARN: yellow,
        LoggingLevels.ERROR: red,
        LoggingLevels.CRITICAL: bold_red,
    }

    def format(self, log_record: logging.LogRecord) -> str:
        """Convert log record into a formatted string.

        Args:
            log_record: LogRecord generated by logging module

        Returns:
            A string formatted according to specifications.
        """
        log_fmt = (
            self.COLORS[LoggingLevels(log_record.levelno)]
            + self.format_template
            + self.reset
        )

        formatter = logging.Formatter(log_fmt)
        formatted_record = formatter.format(log_record)
        return formatted_record


def get_logging_level() -> LoggingLevels:
    """Get logging level from ENV.

    Returns:
        current verbosity of logger

    Raises:
        KeyError: If the verbosity stored in env / constants is not a member
            of the LoggingLevels enum
    """
    verbosity = COML_LOGGING_VERBOSITY.upper()
    if verbosity not in LoggingLevels.__members__:
        raise KeyError(
            f"Verbosity must be one of "
            f"{list(LoggingLevels.__members__.keys())} not {verbosity}"
        )
    return LoggingLevels[verbosity]


def get_console_handler() -> Any:
    """Get console handler for logging.

    Returns:
        stream handler for rich console
    """
    from coalescenceml.constants import console

    console_handler = RichHandler(
        show_path=False,
        omit_repeated_times=False,
        console=console,
        rich_tracebacks=True,
    )
    console_handler.setFormatter(CustomFormatter())
    return console_handler


def get_file_handler() -> Any:
    """Return a file handler for logging.

    Returns:
        file handler for logs
    """
    file_handler = TimedRotatingFileHandler(LOG_FILE, when="midnight")
    file_handler.setFormatter(CustomFormatter())
    return file_handler


def get_logger(logger_name: str) -> logging.Logger:
    """Fetch logger with name.

    Args:
        logger_name: Name of logger to initialize.

    Returns:
        A logger object.
    """
    logger = logging.getLogger(logger_name)
    logger.setLevel(get_logging_level().value)
    logger.addHandler(get_console_handler())

    # TODO: We hand persistent file handing and we need to double
    # logger.addHandler(get_file_handler())
    # check on the propagate issue whether this is OK
    logger.propagate = False
    return logger


def set_root_verbosity() -> None:
    """Set main verbosity of module."""
    level = get_logging_level()
    if level == LoggingLevels.NOTSET:
        logging.getLogger().disabled = True
        get_logger(__name__).debug("Logging NOTSET")
    else:
        rich_tb_install(show_locals=(level == LoggingLevels.DEBUG))

        logging.basicConfig(level=level.value)
        get_logger(__name__).debug(
            f"Logging set to level: {logging.getLevelName(level.value)}"
        )


def init_logging() -> None:
    """Initialize logging with default configuration."""
    set_root_verbosity()

    muted_loggers: List[str] = []
    for logger_name in muted_loggers:
        logging.getLogger(logger_name).setLevel(logging.WARNING)
        logging.getLogger(logger_name).disabled = True

    # ML-Metadata and ml-pipelines-sdk apparently use absl (which makes sense)
    # BUT we want to hide their logging so we will try this.
    from absl import logging as absl_logging

    absl_logging.set_verbosity(-10)
