"""Functions for working with dataclasses.
"""
from __future__ import annotations

import dataclasses as dc
from dataclasses import Field
from typing import Any
from typing import cast
from typing import ClassVar
from typing import Hashable
from typing import Protocol
from typing import Tuple

from pydantic import ConfigDict


pydc_config = ConfigDict(arbitrary_types_allowed=True)


class IsDataclass(Protocol):
    # Checking for the "__dataclass_fields__" attribute is the most reliable way to
    # ascertain that an object is a dataclass, see:
    # https://stackoverflow.com/a/55240861
    # https://github.com/python/typeshed/blob/main/stdlib/_typeshed/__init__.pyi#L353
    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]


class Converter(Protocol):
    """Function signature for a dataclass field value conversion function."""

    def __call__(self, field_value: Any, field_type: str) -> Any:
        ...


def get_args_dict(
    dc_instance: IsDataclass, /, cmp_exclude: list[str] | None = None,
    converter: Converter | None = None,
) -> dict[str, Any]:
    # Custom dict factory function for dataclasses, see:
    # https://github.com/python/cpython/blob/main/Lib/dataclasses.py#L1294
    def dataset_dict_factory(lst: list[tuple[str, Any]]) -> dict[str, Any]:
        new_lst: list[tuple[str, Any]] = []
        for field_name, field_value in lst:
            if cmp_exclude is None or field_name not in cmp_exclude:
                new_lst.append((field_name, field_value))
        return dict(new_lst)

    return _asdict_inner(
        dc_instance, dict_factory=dataset_dict_factory, converter=converter,
    )


def _asdict_inner(obj: Any, dict_factory: Any, converter: Converter | None) -> Any:
    """
    Custom implementation of _asdict_inner from dataclasses.py that avoids deepcopy'ing
    and returns additional information for deserialization. Also supports applying a
    custom conversion function to the dataclass field values.
    """
    if dc._is_dataclass_instance(obj):  # type: ignore[attr-defined]
        fields = []
        for f in dc.fields(obj):
            value = getattr(obj, f.name)
            if converter:
                value = converter(value, str(f.type))
            value = _asdict_inner(value, dict_factory, converter)
            fields.append((f.name, value))
        return {
            'module': str(obj.__module__),
            'class': obj.__class__.__name__,
            'args': dict_factory(fields),
        }
    elif isinstance(obj, (list, tuple)):
        return type(obj)(_asdict_inner(v, dict_factory, converter) for v in obj)
    elif isinstance(obj, dict):
        return type(obj)(
            (
                _asdict_inner(k, dict_factory, converter),
                _asdict_inner(v, dict_factory, converter),
            )
            for k, v in obj.items()
        )
    else:
        # dataclasses.py: return copy.deepcopy(obj)
        # If obj is not a dataclass/list/tuple/dict then return the object itself
        # rather than a deepcopy, as dataclasses.py does. This enables "_asdict_inner"
        # to work on dataclasses that contain objects that cannot be pickled.
        return obj


def get_args_tuple(
    dc_instance: IsDataclass, /, cmp_exclude: list[str] | None = None,
) -> tuple[Hashable]:
    # Helper function to return a non-mutable / hashable repr. of an object
    def make_hashable(value: Any) -> Hashable | tuple[Hashable]:
        if isinstance(value, (list, set, tuple)):
            value = tuple(make_hashable(v) for v in value)
        elif isinstance(value, dict):
            # "hashable" repr. of a dict: https://stackoverflow.com/a/2704866
            dict_keys = value.keys()
            dict_values = [make_hashable(v) for v in value.values()]
            value = tuple(sorted(zip(dict_keys, dict_values)))
        return value

    # NB: we can't use the dataclass astuple() function for this task since it
    # does not provide the dataclass fields to the user-specified tuple factory
    # function, hence there is no way to implement exclusion_fields using this.
    args_dict = get_args_dict(dc_instance, cmp_exclude=cmp_exclude)
    return cast(Tuple[Hashable], make_hashable(args_dict))
