import enum
from typing import Dict, Type, Union, Any, TypeVar
import dataclasses_json.cfg


_EnumT = TypeVar("_EnumT", bound=Union[enum.Enum, enum.IntFlag])


def _enum_entry_decoder(enum_type: Type[_EnumT], entry: Union[int, str, Dict[str, Any]]):
    if isinstance(entry, int):
        return enum_type(entry)
    if isinstance(entry, str):
        return enum_type[entry]
    if isinstance(entry, object):  # pyright: ignore[reportUnnecessaryIsInstance]
        if "name" in entry:
            return enum_type[entry["name"]]
        if "value" in entry:
            enum_type(entry["value"])
    raise ValueError(f"failed to parse enum of type [{enum_type}] from entry [{entry}]")


# NOTE (swapnil): This kinda sucks. Currently, the dataclasses_json class doesn't have a simple way
# to change how enums get encoded and decoded. Ideally, we'd be transmitting name and value for
# enums to make the jsons more human readable.
#
# Stuff that doesn't work:
# - added to_json/from_json/to_dict/from_dict overloads to a new "base enum" class (that inherits
#   from enun.Enum) doesn't work
# - registering enum.Enum to the global encoders/decoders doesn't work, because dataclasses_json
#   isn't doing an isinstance check, its looking up is TYPE is in some DICT. No concept of derived
#   classes in that lookup
# - you can manually specify encoders/decoders every single time an enum is used in a dataclass,
#   but that's error prone and crappy. Also a pain if the enum is nested inside a collection
# - extended the dataclasses_json mixin class would probably work if we did it really carefully, but
#   that feels very hacky and brittle if the library evolves over time. It's possible we silently
#   break things as well.
#
# So with all that, I'm taking this approach, which uses their documented way of overriding
# encoders/decoders for select classes. Add this as a decorator to any enum class being used for
# json encoding/decoding to handle more robust input/output


def patch_json_dataclass_enum_encoder(cls: Type[_EnumT]) -> Type[_EnumT]:
    dataclasses_json.cfg.global_config.encoders[cls] = lambda x: {
        "name": x.name,
        "value": x.value,
    }
    dataclasses_json.cfg.global_config.decoders[cls] = lambda x: _enum_entry_decoder(cls, x)
    return cls
