"""Simple, extendable mechanism for encoding pracitaclly anything to string.

Just register a new encoder for a given type like so:

import draccus
import numpy as np
@draccus.encode.register
def encode_ndarray(obj: np.ndarray) -> str:
    return obj.tostring()
"""
import json
import typing
from argparse import Namespace
from collections.abc import Mapping
from dataclasses import fields, is_dataclass
from enum import Enum
from logging import getLogger
from os import PathLike
from pathlib import PosixPath, WindowsPath
from typing import Any, Dict, Hashable, List, Optional, Tuple, Type, Union

from draccus.choice_types import CHOICE_TYPE_KEY
from draccus.parsers.registry_utils import RegistryFunc, withregistry
from draccus.utils import is_choice_type

logger = getLogger(__name__)


class SimpleJsonEncoder(json.JSONEncoder):
    def default(self, o: Any) -> Any:
        return encode(o)


"""
# NOTE: This code is commented because of static typing check error.
# The problem is incompatibility of mypy and singledispatch.
# See mypy issues for more info:
# https://github.com/python/mypy/issues/8356
# https://github.com/python/mypy/issues/2904
# https://github.com/python/mypy/issues/9112#issuecomment-725316936

class Dataclass(Protocol):
    # see dataclasses.is_dataclass implementation with _FIELDS
    __dataclass_fields__: Dict[str, Field[Any]]


T = TypeVar("T", bool, int, None, str)


@overload
def encode(obj: Dataclass) -> Dict[str, Any]: ...

@overload
def encode(obj: Union[List[Any], Set[Any], Tuple[Any, ...]]) -> List[Any]:
    ...

@overload
def encode(obj: Mapping[Any, Any]) -> Dict[Any, Any]: ...

@overload
def encode(obj: T) -> T: ...
"""


@withregistry
def encode(obj: Any, declared_type: Optional[Type] = None) -> Any:
    """Encodes an object into a json/yaml-compatible primitive type.

    Args:
        obj: The object to encode
        declared_type: Optional type annotation for the object. Used to determine if the object
                      should be encoded as a choice type based on its declared type rather than
                      its concrete type.
    """
    if declared_type is not None:
        underlying_type = typing.get_origin(declared_type) or declared_type
        # we have to handle unions specially for declared types:
        if underlying_type is Union:
            # find the first type that matches the object's type
            for t in typing.get_args(declared_type):
                # we can't use subscripted generic types here
                if typing.get_origin(t) is typing.Literal:
                    for arg in typing.get_args(t):
                        if arg == obj:
                            underlying_type = t
                            declared_type = t
                            break
                elif isinstance(obj, typing.get_origin(t) or t):
                    underlying_type = typing.get_origin(t) or t
                    declared_type = t
                    break
    else:
        underlying_type = type(obj)
    cached_func: RegistryFunc = encode.dispatch(underlying_type)

    if cached_func is None:
        # see if the actual type has a custom encoder
        cached_func = encode.dispatch(type(obj))

    if cached_func is not None:
        fn = cached_func.func

        # we want to support the old interface where the decoding function
        # takes only one argument, so we wrap it here
        try:
            return fn(obj, declared_type)
        except TypeError:
            try:
                return fn(obj)
            except Exception as e:  # pylint: disable=broad-except
                raise Exception(f"Couldn't encode {obj}") from e

    try:
        if underlying_type is not None and is_choice_type(underlying_type):
            return encode_choice(obj, underlying_type)
        elif is_dataclass(obj):
            return encode_dataclass(obj, declared_type)
        elif obj is None:
            return None
        else:
            raise Exception(f"No parser for object {obj} of type {type(obj)}, consider using draccus.encode.register")
    except Exception as e:
        logger.debug(f"Cannot encode object {obj}: {e}")
        raise e


def encode_dataclass(obj: Any, declared_type: Optional[Type] = None):
    d: Dict[str, Any] = dict()

    # Handle type parameters if declared_type is provided
    type_map = {}
    if declared_type is not None:
        origin = typing.get_origin(declared_type)
        if origin is not None and hasattr(origin, "__parameters__"):
            type_args = typing.get_args(declared_type)
            type_vars = origin.__parameters__
            type_map = dict(zip(type_vars, type_args))
        else:
            origin = declared_type

    for field in fields(obj):
        value = getattr(obj, field.name)
        try:
            # If we have a type map, use it to resolve the field's type
            field_type = field.type
            if type_map and hasattr(field_type, "__parameters__"):
                field_type = field_type[tuple(type_map.get(p, p) for p in field_type.__parameters__)]
            d[field.name] = encode(value, field_type)
        except TypeError as e:
            logger.error(f"Unable to encode field {field.name}: {e}")
            raise e
    return d


def encode_choice(obj: Any, declared_type: Type) -> Dict[str, Any]:
    """Encodes an object as a choice type based on its declared type.

    Args:
        obj: The object to encode
        declared_type: The type annotation for the object, which must be a choice type
    """
    if not is_choice_type(declared_type):
        raise ValueError(f"Expected a choice type, got {declared_type}")

    encoded = encode_dataclass(obj, declared_type)

    if not isinstance(encoded, dict):
        raise Exception(f"Choice Class {obj} is not encoded as a dict: {encoded}")

    # Get the choice name from the declared type
    choice_name = obj.get_choice_name(type(obj))
    encoded = {CHOICE_TYPE_KEY: choice_name, **encoded}

    return encoded


@encode.register(Mapping, include_subclasses=True)
def encode_dict(
    obj: Mapping, declared_type: Optional[Type] = None
) -> Union[typing.Mapping[Any, Any], List[Tuple[Any, Any]]]:
    constructor = type(obj)
    result: Union[Mapping, List[Tuple[Any, Any]]] = constructor()

    # Handle type parameters if declared_type is provided
    key_type = None
    value_type = None
    if declared_type is not None:
        type_args = typing.get_args(declared_type)
        if len(type_args) >= 2:
            key_type, value_type = type_args[:2]

    for k, v in obj.items():
        k_ = encode(k, key_type)
        v_ = encode(v, value_type)
        if isinstance(result, list):
            result.append((k_, v_))
        elif isinstance(k_, Hashable):
            result[k_] = v_  # type: ignore
        else:
            # If the encoded key isn't "Hashable", then we store it as a list of tuples
            if isinstance(result, Mapping):
                result = list(result.items())
            result.append((k_, v_))
    return result


@encode.register(Enum, include_subclasses=True)
def encode_enum(obj: Enum, declared_type: Optional[Type] = None) -> str:
    return obj.name


for t in [str, float, int, bool, bytes]:
    encode.register(t, lambda x, _=None: x)


@encode.register(list)
def encode_list(obj: list, declared_type: Optional[Type] = None) -> list:
    item_type = None
    if declared_type is not None:
        type_args = typing.get_args(declared_type)
        if type_args:
            item_type = type_args[0]
    return [encode(x, item_type) for x in obj]


@encode.register(tuple)
def encode_tuple(obj: tuple, declared_type: Optional[Type] = None) -> list:
    if declared_type is not None:
        type_args = typing.get_args(declared_type)
        if type_args and len(type_args) == len(obj):
            return [encode(x, t) for x, t in zip(obj, type_args)]
        elif type_args and len(type_args) == 2 and type_args[1] is Ellipsis:
            item_type = type_args[0]
            return [encode(x, item_type) for x in obj]
    return [encode(x) for x in obj]


@encode.register(set)
def encode_set(obj: set, declared_type: Optional[Type] = None) -> list:
    item_type = None
    if declared_type is not None:
        type_args = typing.get_args(declared_type)
        if type_args:
            item_type = type_args[0]
    return [encode(x, item_type) for x in obj]


encode.register(PathLike, lambda x, _=None: x.__fspath__())
encode.register(PosixPath, lambda x, _=None: x.__fspath__())
encode.register(WindowsPath, lambda x, _=None: x.__fspath__())

encode.register(Namespace, lambda x, _=None: encode(vars(x)))
