import json
import traceback
from datetime import datetime
from inspect import istraceback
from typing import Callable, Iterable, Union, Dict, Optional, List
from datetime import timezone

from aiologger.formatters.base import Formatter
from aiologger.levels import LEVEL_TO_NAME
from aiologger.records import LogRecord
from aiologger.utils import CallableWrapper


LOGGED_AT_FIELDNAME = "logged_at"
LINE_NUMBER_FIELDNAME = "line_number"
FUNCTION_NAME_FIELDNAME = "function"
LOG_LEVEL_FIELDNAME = "level"
MSG_FIELDNAME = "msg"
FILE_PATH_FIELDNAME = "file_path"


class JsonFormatter(Formatter):
    def __init__(
        self,
        serializer: Callable[..., str] = json.dumps,
        default_msg_fieldname: str = None,
    ) -> None:
        super(JsonFormatter, self).__init__()
        self.serializer = serializer
        self.default_msg_fieldname = default_msg_fieldname or MSG_FIELDNAME

    def _default_handler(self, obj):
        if isinstance(obj, datetime):
            return obj.isoformat()
        elif istraceback(obj):
            tb = "".join(traceback.format_tb(obj))
            return tb.strip().split("\n")
        elif isinstance(obj, Exception):
            return "Exception: %s" % repr(obj)
        elif type(obj) is type:
            return str(obj)
        elif isinstance(obj, CallableWrapper):
            return obj()
        return str(obj)

    def format(self, record: LogRecord) -> str:
        """
        Formats a record and serializes it as a JSON str. If record message isnt
        already a dict, initializes a new dict and uses `default_msg_fieldname`
        as a key as the record msg as the value.
        """
        msg: Union[str, dict] = record.msg
        if not isinstance(msg, dict):
            msg = {self.default_msg_fieldname: msg}

        if record.exc_info:
            msg["exc_info"] = record.exc_info
        if record.exc_text:
            msg["exc_text"] = record.exc_text

        return self.serializer(msg, default=self._default_handler)

    @classmethod
    def format_error_msg(cls, record: LogRecord, exception: Exception) -> Dict:
        traceback_info: Optional[List[str]]
        if exception.__traceback__:
            traceback_info = cls.format_traceback(exception.__traceback__)
        else:
            traceback_info = None
        return {
            "record": {
                LINE_NUMBER_FIELDNAME: record.lineno,
                LOG_LEVEL_FIELDNAME: record.levelname,
                FILE_PATH_FIELDNAME: record.filename,
                FUNCTION_NAME_FIELDNAME: record.funcName,
                MSG_FIELDNAME: str(record.msg),
            },
            LOGGED_AT_FIELDNAME: datetime.utcnow().isoformat(),
            "logger_exception": {
                "type": str(type(exception)),
                "exc": str(exception),
                "traceback": traceback_info,
            },
        }


class ExtendedJsonFormatter(JsonFormatter):
    level_to_name_mapping = LEVEL_TO_NAME
    default_fields = frozenset(
        [
            LOG_LEVEL_FIELDNAME,
            LOGGED_AT_FIELDNAME,
            LINE_NUMBER_FIELDNAME,
            FUNCTION_NAME_FIELDNAME,
            FILE_PATH_FIELDNAME,
        ]
    )

    def __init__(
        self,
        serializer: Callable[..., str] = json.dumps,
        default_msg_fieldname: str = None,
        exclude_fields: Iterable[str] = None,
        tz: timezone = None,
    ) -> None:

        super(ExtendedJsonFormatter, self).__init__(
            serializer=serializer, default_msg_fieldname=default_msg_fieldname
        )
        self.tz = tz
        if exclude_fields is None:
            self.log_fields = self.default_fields
        else:
            self.log_fields = self.default_fields - set(exclude_fields)

    def formatter_fields_for_record(self, record: LogRecord):
        """
        :type record: aiologger.records.ExtendedLogRecord
        """
        datetime_serialized = (
            datetime.now(timezone.utc).astimezone(self.tz).isoformat()
        )

        default_fields = (
            (LOGGED_AT_FIELDNAME, datetime_serialized),
            (LINE_NUMBER_FIELDNAME, record.lineno),
            (FUNCTION_NAME_FIELDNAME, record.funcName),
            (LOG_LEVEL_FIELDNAME, self.level_to_name_mapping[record.levelno]),
            (FILE_PATH_FIELDNAME, record.pathname),
        )

        for field, value in default_fields:
            if field in self.log_fields:
                yield field, value

    def format(self, record) -> str:
        """
        :type record: aiologger.records.ExtendedLogRecord
        """
        msg = dict(self.formatter_fields_for_record(record))
        if record.flatten and isinstance(record.msg, dict):
            msg.update(record.msg)
        else:
            msg[MSG_FIELDNAME] = record.msg

        if record.extra:
            msg.update(record.extra)
        if record.exc_info:
            msg["exc_info"] = record.exc_info
        if record.exc_text:
            msg["exc_text"] = record.exc_text

        return self.serializer(
            msg, default=self._default_handler, **record.serializer_kwargs
        )