import logging.config
import os
import threading
import traceback
from typing import TYPE_CHECKING, Any, Dict, Optional, cast

import click
import orjson
from tabulate import tabulate

from owslib.dictutil import flatten
from owslib.enums import Environment
from owslib.functional import traverse
from owslib.logger.constants import LOG_TAG, LOG_TYPE
from owslib.logger.utils import format_timestamp, get_extra_fields

if TYPE_CHECKING:
    from logging import _SysExcInfoType  # noqa

    def get_dd_correlation_context() -> Dict[str, str]:
        return {}

else:
    try:
        # NB! ddtrace breaks mypy
        from ddtrace import tracer

        get_dd_correlation_context = tracer.get_log_correlation_context

    except ImportError:

        def get_dd_correlation_context() -> Dict[str, str]:
            return {}


class DDJsonFormatter(logging.Formatter):
    def __init__(
        self,
        environment: Environment,
        service_name: str,
        service_version: str,
        logger_name: Optional[str] = None,
        pretty: bool = False,
    ) -> None:
        super().__init__()
        self.environment = environment
        self.service_name = service_name
        self.service_version = service_version
        self.logger_name = logger_name

        option = (
            orjson.OPT_SERIALIZE_NUMPY
            | orjson.OPT_SERIALIZE_DATACLASS
            | orjson.OPT_SERIALIZE_UUID
        )

        if pretty:
            option |= orjson.OPT_INDENT_2

        self.option = option

    def serialize(self, message: Dict[str, Any]) -> str:
        return orjson.dumps(message, option=self.option).decode()

    @classmethod
    def format_exception(cls, exc_info: Optional["_SysExcInfoType"]) -> str:
        return "".join(traceback.format_exception(*exc_info)) if exc_info else ""

    def get_dd_correlation_context(self) -> Dict[str, Any]:
        dd_correlation_context = get_dd_correlation_context()
        if not dd_correlation_context.get("env"):
            dd_correlation_context["env"] = self.environment
        if not dd_correlation_context.get("service"):
            dd_correlation_context["service"] = self.service_name
        if not dd_correlation_context.get("version"):
            dd_correlation_context["version"] = self.service_version
        return dd_correlation_context

    def format(self, record: logging.LogRecord) -> str:
        message = {
            # https://docs.datadoghq.com/logs/log_configuration/attributes_naming_convention/
            "message": record.getMessage(),
            "timestamp": format_timestamp(record.created),
            "level": record.levelname,
            "service": self.service_name,
            "environment": self.environment,
            "thread_id": threading.get_ident(),
            "process_id": os.getpid(),
            "log_type": LOG_TYPE,
            "logger_version": "owslib/logger",
            "tag": LOG_TAG,
            # https://docs.datadoghq.com/getting_started/tagging/
            "tags": [f"log_type:{LOG_TYPE}"],
            # DataDog fields
            "dd": self.get_dd_correlation_context(),
            # Logger fields
            "logger": {
                "name": self.logger_name or record.name,
                "pathname": record.pathname,
                "file_name": record.filename,
                "method_name": record.funcName,
                "lineno": record.lineno,
                "thread_name": record.threadName,
            },
        }

        # Add extra fields
        message.update(get_extra_fields(record))

        # If exception, add debug info
        if record.exc_info:
            message.update(
                {
                    "exc_info": self.format_exception(record.exc_info),
                    "stack_info": record.stack_info,
                }
            )

        return self.serialize(message)


class DebugFormatter(logging.Formatter):
    log_colors = {
        "DEBUG": "white",
        "INFO": "green",
        "WARNING": "yellow",
        "ERROR": "red",
        "CRITICAL": "bold_red",
    }

    def __init__(self, debug_extra: bool = False):
        super().__init__()
        self.debug_extra = debug_extra

    def formatMessage(self, record: logging.LogRecord) -> str:
        log_color = self.log_colors.get(record.levelname)
        created = format_timestamp(record.created)
        message = (
            f"{click.style(f'[{created}]', bold=True)} "
            f"{click.style(f'({record.name})', fg='white')} "
            f"{click.style(record.levelname, bold=True, fg=log_color)} - "
            f"{click.style(record.getMessage(), fg=log_color)}"
        )
        if not self.debug_extra:
            return message

        extra = get_extra_fields(record)
        if extra:
            extra_flatten = flatten(
                cast(
                    Dict[Any, Any],
                    traverse(
                        extra,
                        lambda value, path: str(value) if value is not None else value,
                    ),
                )
            )
            extra_msg = tabulate(sorted(extra_flatten.items()), tablefmt="simple")
            message += f"\n{extra_msg}"
        return message