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

from audience_common import __version__
from audience_common.encoders import JSONEncoder
from audience_common.logger.constants import LOG_TAG, LOG_TYPE
from audience_common.logger.utils import get_extra_fields
from audience_common.utils.dictutil import flatten
from audience_common.utils.functional import traverse

try:
    import orjson  # noqa

    ORJSON_INSTALLED = True
except ImportError:
    ORJSON_INSTALLED = False

try:
    import click  # noqa

    CLICK_INSTALLED = True
except ImportError:
    CLICK_INSTALLED = False


try:
    from tabulate import tabulate  # noqa

    TABULATE_INSTALLED = True
except ImportError:
    TABULATE_INSTALLED = False


if TYPE_CHECKING:
    from logging import _SysExcInfoType  # noqa


try:
    from ddtrace import tracer  # noqa

    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: str,
        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
        self.pretty = pretty

    def serialize(self, message: Dict[str, Any]) -> str:
        if ORJSON_INSTALLED:
            option = (
                orjson.OPT_SERIALIZE_NUMPY
                | orjson.OPT_SERIALIZE_DATACLASS
                | orjson.OPT_SERIALIZE_UUID
                | orjson.OPT_UTC_Z
            )
            if self.pretty:
                option |= orjson.OPT_INDENT_2
            return orjson.dumps(message, option=option).decode()
        return json.dumps(message, cls=JSONEncoder)

    @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, str]:
        dd_correlation_context = {
            **get_dd_correlation_context(),
            "version": self.service_version,
        }
        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
        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": datetime.utcfromtimestamp(record.created).isoformat(),
            "level": record.levelname,
            "levelname": record.levelname,
            "service": self.service_name,
            "environment": self.environment,
            "thread_id": threading.get_ident(),
            "process_id": os.getpid(),
            "log_type": LOG_TYPE,
            "logger_name": self.logger_name or record.name,
            "logger_version": f"audience-common/logger/{__version__}",
            "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 = datetime.utcfromtimestamp(record.created).isoformat()
        if CLICK_INSTALLED:
            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)}"
            )
        else:
            message = (
                f"[{created}] ({record.name}) {record.levelname}"
                f" - {record.getMessage()}"
            )
        if not self.debug_extra:
            return message

        extra = get_extra_fields(record)
        if extra and TABULATE_INSTALLED:
            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