import os
from logging import (
    Formatter,
    getLogger,
    StreamHandler
)
from uuid import uuid4

from owslogger.logger import (
    DSNHandler,
    OwsLoggingAdapter
)

from config import (
    LOGGER_LEVEL,
    APPLICATION_NAME,
    LOGGER_DSN,
    ENVIRONMENT,
    APP_VERSION
)


class ColorFormatter(Formatter):
    """Custom formatter to add colors to log levels."""

    COLORS = {
        'DEBUG': '\033[94m',  # Blue
        'INFO': '\033[92m',   # Green
        'WARNING': '\033[93m',  # Yellow
        'ERROR': '\033[91m',  # Red
        'CRITICAL': '\033[95m'  # Magenta
    }
    RESET = '\033[0m'

    def format(self, record):
        log_color = self.COLORS.get(record.levelname, self.RESET)
        record.levelname = f"{log_color}{record.levelname}{self.RESET}"
        return super().format(record)


class ConcurrencyDSNHandler(DSNHandler):
    """Custom DSN handler.

    Custom DSN handler that sends ensures that the correlation ID is passed
    through to the provider.
    """

    def __init__(self, dsn, logger_level, environment, service_name,
                 service_version, correlation_id=None):
        """DSN Handler for Orch infra. Formats logs for DataDog processing.

        Args:
            dsn (str): the full dsn to send the log to.
            logger_level (str): the logger level.
            environment (str): the current environment.
            service_name (str): the service name.
            service_version (str): the service version.
            correlation_id (str, optional): the correlation ID. Defaults to
                None.

        Returns:
            DSNHandler: the created handler.

        """
        DSNHandler.__init__(
            self,
            dsn,
            environment,
            service_name,
            service_version
        )
        if correlation_id:
            self.correlation_id = correlation_id
        else:
            self.correlation_id = uuid4()

        self.formatter = ColorFormatter(
            '%(asctime)s | %(levelname)s | '
            '%(filename)s:%(funcName)s:%(lineno)d | '
            '%(message)s',
            datefmt='%Y-%m-%d at %I:%M:%S %p %Z'
        )
        self.setLevel(logger_level)

    def emit(self, record):
        """Emit a record.

        From the documentation: do whatever it takes to actually log the
        specified logging record. Here: we send it to the provider. The payload
        matches Orchard format.

        Args:
            record (LogRecord): the record to log.

        """
        record.correlation_id = self.correlation_id
        super().emit(record)


log = getLogger(__name__)
log.setLevel(LOGGER_LEVEL)

formatter = ColorFormatter(
    '%(asctime)s | %(levelname)s | %(filename)s:%(funcName)s:%(lineno)d | '
    '%(message)s',
    datefmt='%Y-%m-%d at %I:%M:%S %p %Z'
)

handler = StreamHandler()
handler.setLevel(LOGGER_LEVEL)
handler.setFormatter(formatter)

log.addHandler(handler)

# PropogateHandler for ows-logger
if LOGGER_DSN:
    corr_id = os.environ.get('CORRELATION_ID', None)
    if not corr_id:
        corr_id = str(uuid4())
        os.environ['CORRELATION_ID'] = corr_id
        log.debug(f"Correlation ID: {corr_id}")

    custom_handler = ConcurrencyDSNHandler(
        LOGGER_DSN, LOGGER_LEVEL, ENVIRONMENT, APPLICATION_NAME, APP_VERSION,
        corr_id)
    custom_handler.setFormatter(formatter)

    # bind context vars for datadog
    log.addHandler(custom_handler)

    context = dict(correlation_id=corr_id)
    log = OwsLoggingAdapter(log, context)