import sys

from collections import OrderedDict
from datetime import datetime
from typing import Any, Literal, List

from loguru import logger

from pretty_json_loguru.get_traceback import get_traceback

try:
    from loguru import Record
except ImportError:
    # for some reason, Record does not import this way in loguru 0.6.0
    Record = Any

try:
    import ujson as json
except ImportError:
    import json

BUILTIN_KEYS = [
    "ts",
    "module",
    "msg",
    "source",
    "error",
    "traceback",
    "level",
]  # note: extra is not included, it is a placeholder for extra fields


def pretty_json_loguru_formatter(
    colorize: bool = True,
    traceback: Literal["attach", "extra", "drop"] = "attach",
    keys: List[
        Literal["ts", "module", "msg", "source", "extra", "error", "traceback", "level"]
    ] = [
        "ts",
        # "module", # module is skipped by default for brevity
        "msg",
        "source",
        "extra",
        "error",
        "traceback",
        "level",
    ],
):
    """Loguru formatter builder for colored json logs.

    Sample output (colored in the console):
    {"ts": "2024-07-29 08:19:03.675", "module": "format_as_colored_json", "message": "Simple message"}
    {"ts": "2024-07-29 08:19:03.675", "module": "format_as_colored_json", "message": "Message with extra", "foo": "bar"}
    {"ts": "2024-07-29 08:19:03.675", "module": "format_as_colored_json", "message": "Exception caught", "error": "ValueError: This is an exception", "traceback": "...\nValueError: This is an exception"}

    Parameters
    ----------
    colorize : bool
        If True, colors will be added to the log. If colorize=False, vanilla traceback will be used (`traceback.format_exc()`)
    traceback : Literal["attach", "extra", "drop"]
        If "attach", traceback will be appended to the log, as if we use the vanilla formatter.
        if "extra", traceback will be added to the extra field
        if "drop", traceback will be dropped
    keys : List[Literal["ts", "module", "msg", "source", "extra", "error", "traceback", "level"]]
        List and order of keys to include in the log. `extra` is a placeholder for extra fields

    """

    # - Check traceback value

    if traceback not in ["attach", "extra", "drop"]:
        raise ValueError(f"Unknown traceback value {traceback}")

    # - Define output formatter

    def _format_as_json_colored(record: Record):
        """
        record:
            {
              "elapsed": "0:00:00.005652",
              "exception": [
                "<class 'ValueError'>",
                "This is an exception",
                "<traceback object at 0x101444c00>"
              ],
              "extra": {
                "foo": "bar"
              },
              "file": "(name='format_as_colored_json.py', path='/Users/marklidenberg/Documents/coding/repos/marklidenberg/pretty-json-loguru/pretty_json_loguru/formatters/format_as_colored_json.py')",
              "function": "test",
              "level": "(name='ERROR', no=40, icon='\u274c')",
              "line": 225,
              "message": "Exception caught",
              "module": "format_as_colored_json",
              "name": "__main__",
              "process": "(id=96919, name='MainProcess')",
              "thread": "(id=8532785856, name='MainThread')",
              "time": "2025-05-09 11:18:46.707576+02:00"
            }
        """

        # - Pop extra

        extra = dict(record["extra"])
        extra.pop("source", None)

        # - Create record_dic that will be serialized as json

        record_dic = {
            "ts": datetime.fromisoformat(str(record["time"])).strftime(
                "%Y-%m-%d %H:%M:%S.%f"
            )[:-3],  # 2023-03-26 13:04:09.512
            "module": record["module"],
            "msg": record["message"],
            "source": record["extra"].get("source", ""),
        }

        if record["exception"] and traceback == "extra":
            record_dic["traceback"] = get_traceback(
                exception=record["exception"],
                colorize=False,
            ).strip()
            record_dic["error"] = record_dic["traceback"].split("\n")[-1]

        record_dic = {k: v for k, v in record_dic.items() if v}
        record_dic.update(extra)

        """
        {
          "msg": "Exception caught",
          "ts": "2025-05-09 11:18:46.707",
          "traceback": "Traceback (most recent call last):\n  File \"/Users/marklidenberg/Documents/coding/repos/marklidenberg/pretty-json-loguru/pretty_json_loguru/formatters/format_as_colored_json.py\", line 223, in test\n    raise ValueError(\"This is an exception\")\nValueError: This is an exception",
          "error": "ValueError: This is an exception",
          "foo": "bar"
        }
        """

        # - Sort keys

        key_to_index = {key: i for i, key in enumerate(keys)}

        def _key_func(kv):
            if kv[0] in BUILTIN_KEYS:
                # default keys
                if kv[0] in keys:
                    return key_to_index[kv[0]]
                else:
                    return len(keys)
            else:
                # extra keys, e.g. foo=bar
                if "extra" in keys:
                    return keys.index("extra")
                else:
                    return len(keys)

        record_dic = OrderedDict(
            sorted(
                record_dic.items(),
                key=_key_func,
            )
        )

        # - Filter keys

        def _filter_func(k):
            if k in BUILTIN_KEYS:
                return k in keys
            else:
                return "extra" in keys

        record_dic = {k: v for k, v in record_dic.items() if _filter_func(k)}

        # - Get json

        output = (
            json.dumps(
                record_dic,
                default=str,
                ensure_ascii=False,
            )
            .replace("{", "{{")
            .replace(
                "}",
                "}}",
            )
        )

        # - Iterate over json and add color tags

        for i, (key, value) in enumerate(record_dic.items()):
            # - Dump to json

            value_str = (
                json.dumps(value, default=str, ensure_ascii=False)
                .replace("{", "{{")
                .replace("}", "}}")
            )

            if colorize:
                # - Init level colors

                """
                Original colors from loguru:
                | Level        | Default Color Tag |               |
                | ------------ | ----------------- | ------------- |
                | **TRACE**    | `<cyan><bold>`    |               |
                | **DEBUG**    | `<blue><bold>`    |               |
                | **INFO**     | `<bold>`          |               |
                | **SUCCESS**  | `<green><bold>`   |               |
                | **WARNING**  | `<yellow><bold>`  |               |
                | **ERROR**    | `<red><bold>`     |               |
                | **CRITICAL** | `<RED><bold>`     | ([GitHub][1]) |
                
                [1]: https://github.com/Delgan/loguru/blob/master/loguru/_defaults.py "loguru/loguru/_defaults.py at master · Delgan/loguru · GitHub"
                """

                level_colors = {
                    "TRACE": "white",
                    "DEBUG": "blue",
                    "INFO": "light-white",
                    "SUCCESS": "green",
                    "WARNING": "yellow",
                    "ERROR": "red",
                    "CRITICAL": "light-white",
                }

                # - Pick color

                color_key = {
                    "ts": "green",
                    "module": "cyan",
                    "msg": level_colors[record["level"].name],
                }.get(key, "magenta")

                color_value = {
                    "ts": "green",
                    "module": "cyan",
                    "msg": level_colors[record["level"].name],
                }.get(key, "yellow")

                # - Add colors for keys and values

                colored_key = (
                    f'<{color_key}>"{{extra[_extra_{2 * i}]}}"</{color_key}>'
                    if color_key
                    else f'"{{extra[_extra_{2 * i}]}}"'
                )
                colored_value = (
                    f"<{color_value}>{{extra[_extra_{2 * i + 1}]}}</{color_value}>"
                    if color_value
                    else f"{{extra[_extra_{2 * i + 1}]}}"
                )

                if key == "msg" and record["level"].name == "CRITICAL":
                    colored_key = f"<RED>{colored_key}</RED>"
                    colored_value = f"<RED>{colored_value}</RED>"

                output = output.replace(
                    f'"{key}": {value_str}',
                    f"{colored_key}: {colored_value}",
                )

            # - Add key and value to record, from where loguru will get them

            if record:
                record["extra"][f"_extra_{2 * i}"] = key
                record["extra"][f"_extra_{2 * i + 1}"] = json.dumps(
                    value,
                    ensure_ascii=False,
                    default=str,
                )

        # - Add traceback on new line

        if traceback == "attach" and record["exception"]:
            record["extra"]["_extra_traceback"] = get_traceback(
                exception=record["exception"],
                colorize=colorize,
            )
            output += "\n{extra[_extra_traceback]}"

        # - Add white color for the whole output

        return "<white>" + output + "\n" + "</white>"

    return _format_as_json_colored


def test():
    logger.trace("Trace message")
    logger.debug("Debug message")
    logger.info("Info message")
    logger.success("Success message")
    logger.warning("Warning message")
    logger.error("Error message")
    logger.critical("Critical message")

    logger.info("Message with extra", foo="bar")

    try:
        raise ValueError("This is an exception")
    except ValueError:
        logger.exception("Exception caught", foo="bar")


if __name__ == "__main__":
    logger.remove()
    logger.add(
        sys.stdout,
        format=pretty_json_loguru_formatter(
            # colorize=True,
            # traceback="extra",
        ),
        level="TRACE",
    )

    #
    test()
