import os
import logging
from logging.handlers import TimedRotatingFileHandler


class ColorCodes:
    GREY = "\x1b[38;21m"
    BLUE = "\x1b[38;5;39m"
    YELLOW = "\x1b[38;5;226m"
    RED = "\x1b[38;5;196m"
    BOLD_RED = "\x1b[31;1m"
    RESET = "\x1b[0m"


class ColorFormatter(logging.Formatter):
    def __init__(self, fmt: str):
        super().__init__()
        self.fmt = fmt
        self.FORMATS = {
            logging.DEBUG: ColorCodes.GREY + fmt + ColorCodes.RESET,
            logging.INFO: ColorCodes.BLUE + fmt + ColorCodes.RESET,
            logging.WARNING: ColorCodes.YELLOW + fmt + ColorCodes.RESET,
            logging.ERROR: ColorCodes.RED + fmt + ColorCodes.RESET,
            logging.CRITICAL: ColorCodes.BOLD_RED + fmt + ColorCodes.RESET,
        }

    def format(self, record):
        log_fmt = self.FORMATS.get(record.levelno)
        formatter = logging.Formatter(log_fmt)
        return formatter.format(record)


def get_logger(stage: str = "dev", port: int = 50051) -> logging.Logger:
    """
    :param port:
    :param stage:
    :return:
    """
    name = "word_segmentation"
    log_dir = "logs"
    logger = logging.getLogger(name)
    if logger.handlers:
        logger.debug("Logger already configured")
        return logger

    log_format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    if stage == "dev":
        logger.setLevel(logging.DEBUG)
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(ColorFormatter(log_format))
        logger.addHandler(console_handler)
    else:
        logger.setLevel(logging.INFO)
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
        file_handler = TimedRotatingFileHandler(
            filename=os.path.join(log_dir, f"{name}_{port}.log"),
            when="midnight",
            interval=1,
            backupCount=30,
            encoding="utf-8",
        )
        file_handler.setFormatter(logging.Formatter(log_format))
        logger.addHandler(file_handler)
    return logger
