import json
from collections import defaultdict

from rex.utils.io import tensor_friendly_json_encoding


def get_dict_content(dict_item: dict, key: str):
    val = dict_item
    for k in key.split("."):
        val = val[k]
    return val


def flatten_dict(d, parent_key="", sep="."):
    items = []
    for k, v in d.items():
        new_key = parent_key + sep + k if parent_key else k
        if isinstance(v, dict):
            items.extend(flatten_dict(v, new_key, sep=sep).items())
        else:
            items.append((new_key, v))
    return dict(items)


def _pretty_format(
    obj: dict, decimal=3, return_percentage=False, add_percentage_symbol=True
):
    if not isinstance(obj, dict):
        return obj

    pretty = {}
    for key, val in obj.items():
        if isinstance(val, float):
            new_val = val
            if return_percentage:
                new_val = 100 * new_val
            if add_percentage_symbol:
                template = f"{{:.{decimal}f}} %"
            else:
                template = f"{{:.{decimal}f}}"
            new_val = template.format(new_val)
        elif isinstance(val, dict):
            new_val = _pretty_format(
                val,
                decimal=decimal,
                return_percentage=return_percentage,
                add_percentage_symbol=add_percentage_symbol,
            )
        else:
            new_val = val
        pretty[key] = new_val
    return pretty


class PrettyPrintDict(dict):
    def __str__(
        self, indent=2, decimal=3, return_percentage=True, add_percentage_symbol=True
    ):
        pretty_dict = _pretty_format(
            self,
            decimal=decimal,
            return_percentage=return_percentage,
            add_percentage_symbol=add_percentage_symbol,
        )
        return json.dumps(
            pretty_dict,
            ensure_ascii=False,
            indent=indent,
            default=tensor_friendly_json_encoding,
        )

    def to_dict(self) -> dict:
        return dict(self)

    def jsonify(self, **kwargs) -> str:
        return json.dumps(
            dict(self),
            ensure_ascii=False,
            default=tensor_friendly_json_encoding,
            **kwargs,
        )


class PrettyPrintDefaultDict(defaultdict):
    def __str__(
        self, indent=2, decimal=3, return_percentage=True, add_percentage_symbol=True
    ):
        pretty_dict = _pretty_format(
            self,
            decimal=decimal,
            return_percentage=return_percentage,
            add_percentage_symbol=add_percentage_symbol,
        )
        return json.dumps(
            pretty_dict,
            ensure_ascii=False,
            default=tensor_friendly_json_encoding,
            indent=indent,
        )

    def to_dict(self) -> dict:
        return dict(self)

    def jsonify(self, **kwargs) -> str:
        return json.dumps(
            dict(self),
            ensure_ascii=False,
            default=tensor_friendly_json_encoding,
            **kwargs,
        )
