# GT4Py - GridTools Framework
#
# Copyright (c) 2014-2024, ETH Zurich
# All rights reserved.
#
# Please, refer to the LICENSE file in the root directory.
# SPDX-License-Identifier: BSD-3-Clause

from __future__ import annotations

import copyreg
import dataclasses
import functools
import inspect
import itertools
import types
from typing import (
    Any,
    Callable,
    ClassVar,
    Final,
    Optional,
    ParamSpec,
    Sequence,
    TypeAlias,
    TypeGuard,
    TypeVar,
    cast,
    overload,
)

from gt4py.eve import datamodels, utils as eve_utils


GT4PY_CLASS_METADATA_NS: Final[str] = "GT4PY_META"


def gt4py_metadata(**kwargs: Any) -> dict[str, dict[str, Any]]:
    """Helper function to store dataclass/datamodel field metadata within a GT4Py namespace."""
    return {GT4PY_CLASS_METADATA_NS: kwargs}


_StandardPickleState = None | dict[str, Any] | tuple[dict[str, Any] | None, dict[str, Any]]
_StandardGetStateMethod: TypeAlias = Callable[[object], _StandardPickleState]
_StandardSetStateMethod: TypeAlias = Callable[[object, _StandardPickleState], None]


@functools.cache
def _get_metadata_based_state_getstate(cls: type) -> _StandardGetStateMethod:
    """
    Helper function to make class-specific `__getstate__` method following the standard implementation.
    """
    if not isinstance(cls, type) or not (
        (is_dataclass := dataclasses.is_dataclass(cls)) or datamodels.is_datamodel(cls)
    ):
        raise TypeError(f"Expected a dataclass or datamodel type, got '{cls}'")

    if is_dataclass:
        field_metadata = {
            field.name: field.metadata.get(GT4PY_CLASS_METADATA_NS, None)
            for field in dataclasses.fields(cls)
        }
    else:
        field_metadata = {
            field.name: field.metadata.get(GT4PY_CLASS_METADATA_NS, None)
            for field in datamodels.get_fields(cls).values()
        }

    # To gather all slots we need to traverse the whole MRO not just the class itself.
    # We reuse the implementation of `copyreg._slotnames` to avoid code duplication and
    # potential bugs, even if it is not (unfortunately) part of the module's public API.
    class_slots = copyreg._slotnames(cls)  # type: ignore[attr-defined] # copyreg._slotnames is not recognized

    has_slots = len(class_slots) > 0
    has_dict = any("__dict__" in c.__dict__ for c in cls.__mro__)

    dict_names = []
    slot_names = []
    for name, metadata in field_metadata.items():
        if metadata and not metadata.get("pickle", True):
            continue
        if name in class_slots:
            slot_names.append(name)
        else:
            dict_names.append(name)

    # Comply with the default implementation of object.__getstate__() / object.__setstate__(state)
    if not (has_slots or has_dict):
        # for class instances without __dict__ nor __slots__: state = None.
        def __getstate__(self: object) -> _StandardPickleState:
            return None

    elif not has_slots:
        # for class instances with __dict__ and no __slots__: state = self.__dict__.
        def __getstate__(self: object) -> _StandardPickleState:
            return {name: getattr(self, name) for name in dict_names}

    elif not has_dict:
        # for class instances with __slots__ and no __dict__: state = (None, {slot.name: slot.value for all slots}).
        def __getstate__(self: object) -> _StandardPickleState:
            return (None, {name: getattr(self, name) for name in slot_names})

    else:
        # for class instances with __dict__ and __slots__: state = (self.__dict__, {slot.name: slot.value for all slots}).
        def __getstate__(self: object) -> _StandardPickleState:
            return (
                {name: getattr(self, name) for name in dict_names},
                {name: getattr(self, name) for name in slot_names},
            )

    return __getstate__


class MetadataBasedPickling:
    """
    Mixin for adding metadata-based pickling to dataclass-like objects.

    It uses the class field information to select only instance fields which
    are not marked with `pickle=False` in the 'GT4PY_META' metadata namespace.
    Individual fields can therefore opt out of pickling.
    For example: `foo = field(..., metadata=gt4py_metadata(pickle=False))`.
    """

    __slots__ = ()  # to avoid creation of __dict__ when not needed

    def __getstate__(self) -> _StandardPickleState:
        """
        Get the state of the object for pickling.

        It returns the same kind of arguments as the default `__getstate__`
        implementation, used by `pickle` (as documented in `object.__getstate__`,
        check: https://devdocs.io/python~3.14/library/pickle#object.__getstate__)
        """
        return _get_metadata_based_state_getstate(type(self))(self)  # type: ignore[arg-type]  # type(self) should be hashable

    # Note: we don't implement `__setstate__` as the output of our custom
    #  `__getstate__` implementation should be compatible with the default
    #  implementation.


class RecursionGuard:
    """
    Context manager to guard against inifinite recursion.

    >>> def foo(i):
    ...     with RecursionGuard(i):
    ...         if i % 2 == 0:
    ...             foo(i)
    ...     return i
    >>> foo(3)
    3
    >>> foo(2)  # doctest:+ELLIPSIS
    Traceback (most recent call last):
        ...
    gt4py.next.utils.RecursionGuard.RecursionDetected
    """

    guarded_objects: ClassVar[set[int]] = set()

    obj: Any

    class RecursionDetected(Exception):
        pass

    def __init__(self, obj: Any):
        self.obj = obj

    def __enter__(self) -> None:
        if id(self.obj) in self.guarded_objects:
            raise self.RecursionDetected()
        self.guarded_objects.add(id(self.obj))

    def __exit__(self, *exc: Any) -> None:
        self.guarded_objects.remove(id(self.obj))


_T = TypeVar("_T")
_P = ParamSpec("_P")
_R = TypeVar("_R")


def is_tuple_of(v: Any, t: type[_T]) -> TypeGuard[tuple[_T, ...]]:
    return isinstance(v, tuple) and all(isinstance(e, t) for e in v)


# TODO(havogt): remove flatten duplications in the whole codebase
def flatten_nested_tuple(
    value: tuple[
        _T | tuple, ...
    ],  # `_T` omitted on purpose as type of `value`, to properly deduce `_T` on the user-side
) -> tuple[_T, ...]:
    if isinstance(value, tuple):
        return sum((flatten_nested_tuple(v) for v in value), start=())  # type: ignore[arg-type] # cannot properly express nesting
    else:
        return (value,)


@overload
def tree_map(
    fun: Callable[_P, _R],
    *,
    collection_type: type | tuple[type, ...] = tuple,
    result_collection_constructor: Optional[Callable] = None,
    unpack: bool = False,
    with_path_arg: bool = False,
) -> Callable[..., _R | tuple[_R | tuple, ...]]: ...


@overload
def tree_map(
    *,
    collection_type: type | tuple[type, ...] = tuple,
    result_collection_constructor: Optional[Callable] = None,
    unpack: bool = False,
    with_path_arg: bool = False,
) -> Callable[
    [Callable[_P, _R]], Callable[..., Any]
]: ...  # TODO(havogt): typing of `result_collection_constructor` is too weak here


def tree_map(
    fun: Optional[Callable[_P, _R]] = None,
    *,
    collection_type: type | tuple[type, ...] = tuple,
    result_collection_constructor: Optional[Callable] = None,
    unpack: bool = False,
    with_path_arg: bool = False,
) -> Callable[..., _R | tuple[_R | tuple, ...]] | Callable[[Callable[_P, _R]], Callable[..., Any]]:
    """
    Apply `fun` to each entry of (possibly nested) collections (by default `tuple`s).

    Args:
        fun: Function to apply to each entry of the collection.
        collection_type: Type of the collection to be traversed. Can be a single type or a tuple of types.
        result_collection_constructor: Type of the collection to be returned. If `None` the same type as `collection_type` is used.
        unpack: Replicate tuple structure returned from `fun` to the mapped result, i.e. return
          tuple of result collections instead of result collections of tuples.
        with_path_arg: Pass the path to access the current element to `fun`.
    Examples:
        >>> tree_map(lambda x: x + 1)(((1, 2), 3))
        ((2, 3), 4)

        >>> tree_map(lambda x, y: x + y)(((1, 2), 3), ((4, 5), 6))
        ((5, 7), 9)

        >>> tree_map(collection_type=list)(lambda x: x + 1)([[1, 2], 3])
        [[2, 3], 4]

        >>> tree_map(
        ...     collection_type=(list, tuple),
        ...     result_collection_constructor=lambda value, elts: tuple(elts)
        ...     if isinstance(value, list)
        ...     else list(elts),
        ... )(lambda x: x + 1)([(1, 2), 3])
        ([2, 3], 4)

        >>> @tree_map
        ... def impl(x):
        ...     return x + 1
        >>> impl(((1, 2), 3))
        ((2, 3), 4)

        >>> @tree_map(with_path_arg=True)
        ... def impl(x, path: tuple[int, ...]):
        ...     path_str = "".join(f"[{i}]" for i in path)
        ...     return f"t{path_str} = {x}"
        >>> t = impl(((1, 2), 3))
        >>> t[0][0]
        't[0][0] = 1'
        >>> t[0][1]
        't[0][1] = 2'
        >>> t[1]
        't[1] = 3'

        >>> @tree_map(unpack=True)
        ... def impl(x):
        ...     return (x, x**2)
        >>> identity, squared = impl(((2, 3), 4))
        >>> identity
        ((2, 3), 4)
        >>> squared
        ((4, 9), 16)
    """

    if result_collection_constructor is None:
        if isinstance(collection_type, tuple):
            # Note: that doesn't mean `collection_type=tuple`, but e.g. `collection_type=(list, tuple)`
            raise TypeError(
                "tree_map() requires `result_collection_constructor` when `collection_type` is a tuple of types."
            )
        result_collection_constructor = lambda _, elts: collection_type(elts)  # noqa: E731 # because a lambda is clearer

    if fun:

        @functools.wraps(fun)
        def impl(*args: Any | tuple[Any | tuple, ...]) -> _R | tuple[_R | tuple, ...]:
            if isinstance(args[0], collection_type):
                non_path_args: Sequence[Any]
                if with_path_arg:
                    *non_path_args, path = args
                    args = (*non_path_args, tuple((*path, i) for i in range(len(args[0]))))
                else:
                    non_path_args = args

                assert all(
                    isinstance(arg, collection_type) and len(args[0]) == len(arg)
                    for arg in non_path_args
                )
                assert result_collection_constructor is not None
                ctor = functools.partial(result_collection_constructor, args[0])

                mapped = [impl(*arg) for arg in zip(*args)]
                if unpack:
                    return tuple(map(ctor, zip(*mapped)))
                else:
                    return ctor(mapped)

            return fun(  # type: ignore[call-arg]
                *cast(_P.args, args),  # type: ignore[valid-type]
            )  # mypy doesn't understand that `args` at this point is of type `_P.args`

        if with_path_arg:
            return lambda *args: impl(*args, ())
        else:
            return impl
    else:
        return functools.partial(
            tree_map,
            collection_type=collection_type,
            result_collection_constructor=result_collection_constructor,
            unpack=unpack,
            with_path_arg=with_path_arg,
        )


_T1 = TypeVar("_T1")
_T2 = TypeVar("_T2")


def equalize_tuple_structure(
    d1: _T1, d2: _T2, *, fill_value: Any, bidirectional: bool = True
) -> tuple[_T1, _T2]:
    """
    Given two values or tuples thereof equalize their structure.

    If one of the arguments is a tuple the argument will be promoted to a tuple of the same
    structure.

    >>> equalize_tuple_structure((1,), (2, 3), fill_value=None) == (
    ...     (1, None),
    ...     (2, 3),
    ... )
    True

    >>> equalize_tuple_structure(1, (2, 3), fill_value=None) == (
    ...     (1, 1),
    ...     (2, 3),
    ... )
    True

    If `bidirectional` is `False` only the second argument is equalized and an error is raised if
    the first argument would require promotion.
    """
    d1_is_tuple = isinstance(d1, tuple)
    d2_is_tuple = isinstance(d2, tuple)
    if not d1_is_tuple and d2_is_tuple:
        if not bidirectional:
            raise ValueError(f"Expected `{d2!s}` to have same structure as `{d1!s}`.")
        return equalize_tuple_structure(
            (d1,) * len(d2),  # type: ignore[arg-type] # ensured by condition above
            d2,
            fill_value=fill_value,
            bidirectional=bidirectional,
        )
    if d1_is_tuple and not d2_is_tuple:
        return equalize_tuple_structure(
            d1,
            (d2,) * len(d1),  # type: ignore[arg-type] # ensured by condition above
            fill_value=fill_value,
            bidirectional=bidirectional,
        )
    if d1_is_tuple and d2_is_tuple:
        if not bidirectional and len(d1) < len(d2):  # type: ignore[arg-type] # ensured by condition above
            raise ValueError(f"Expected `{d2!s}` to be longer than or equal to `{d1!s}`.")
        return tuple(  # type: ignore[return-value] # mypy not smart enough
            zip(
                *(
                    equalize_tuple_structure(
                        el1, el2, fill_value=fill_value, bidirectional=bidirectional
                    )
                    for el1, el2 in itertools.zip_longest(d1, d2, fillvalue=fill_value)  # type: ignore[call-overload] # d1, d2 are tuples
                )
            )
        )
    return d1, d2


def make_args_canonicalizer(
    signature: inspect.Signature,
    *,
    name: str | None = None,
) -> Callable[..., tuple[tuple, dict[str, Any]]]:
    """
    Create a call arguments canonicalizer function from a given signature.

    The canonicalizer returns the call arguments in a tuple with the
    following two items:

        - a tuple of values containing all the positional parameters
        - a dictionary of parameter names to values for all keyword-only parameters

    Args:
        signature: The signature for which to create the canonicalizer.

    Keyword Args:
        name: Name of the function for which the canonicalizer is created.

    """

    params = []
    pos_args = []
    key_args = []
    var_pos_arg: str | None = None
    var_key_arg: str | None = None

    for key, param in signature.parameters.items():
        match param.kind:
            case inspect.Parameter.POSITIONAL_ONLY | inspect.Parameter.POSITIONAL_OR_KEYWORD:
                pos_args.append(f"{key}")
            case inspect.Parameter.VAR_POSITIONAL:
                var_pos_arg = key
                pos_args.append(f"*{key}")
            case inspect.Parameter.KEYWORD_ONLY:
                key_args.append(f"'{key}': {key}")
            case inspect.Parameter.VAR_KEYWORD:
                var_key_arg = key
                key_args.append(f"**{key}")

        # Remove annotations to avoid issues in case `str(annotation)`
        # produces incorrect python expressions
        params.append(param.replace(annotation=param.empty))

    canonicalizer_signature = inspect.Signature(parameters=params)

    if len(pos_args) == 1:
        if var_pos_arg:
            # If there is only an '*args' parameter, we can just return it directly
            pos_args_tuple_expr = var_pos_arg
        else:
            pos_args_tuple_expr = f"({pos_args[0]},)"  # Add trailing comma to make it a tuple
    else:
        # In the regular case, assemble the output tuple with all positional arguments
        pos_args_tuple_expr = f"({str.join(', ', pos_args)})"

    if len(key_args) == 1 and var_key_arg:
        # If there is only a '**kwargs' parameter, we can just return it directly
        key_args_dict_expr = var_key_arg
    else:
        # In the regular case, assemble the output dictionary with all keyword arguments
        key_args_dict_expr = f"{{ {str.join(', ', key_args)} }}"

    canonicalize_func_name = f"canonicalize_args_for_{name}" if name else "canonicalize_args"
    canonicalizer_src = f"""
def {canonicalize_func_name}{canonicalizer_signature!s}:
    return {pos_args_tuple_expr}, {key_args_dict_expr}
"""
    namespace: dict[str, Any] = {}
    exec(canonicalizer_src, namespace)
    canonicalizer = namespace[canonicalize_func_name]

    return cast(Callable[..., tuple[tuple, dict[str, Any]]], canonicalizer)


@functools.cache
def make_args_canonicalizer_for_function(
    func: types.FunctionType,
) -> Callable[..., tuple[tuple, dict[str, Any]]]:
    return make_args_canonicalizer(inspect.signature(func), name=func.__name__)


def canonicalize_call_args(
    func: Callable,
    /,
    args: tuple,
    kwargs: dict[str, Any],
) -> tuple[tuple, dict[str, Any]]:
    """
    Canonicalize call arguments for a given function.

    Args:
        func: The function for which to canonicalize the call arguments.
        args: Positional arguments.
        kwargs: Keyword arguments.

    Returns:
        A tuple of positional arguments and a dictionary with keyword arguments.

    Note:
        This is a convenience wrapper around `make_args_canonicalizer_for_function`.
    """

    return make_args_canonicalizer_for_function(func)(*args, **kwargs)


class IDGeneratorPool(eve_utils.CustomDefaultDictBase):
    """
    Utility for providing unique IDs for each prefix.

    The use-case for this implementation is to provide a single IDGeneratorPool for the whole
    transformation pipeline, which ensure unique names (prefix+counter) across all transformations.

    The reason for this implementation of a single counter per prefix (over a global increasing counter)
    is that parts of the IR might still be stable when comparing different versions of a transformation.

    >>> uids = IDGeneratorPool()
    >>> next(uids["foo"])
    'foo_0'
    >>> next(uids["foo"])
    'foo_1'
    >>> next(uids["bar"])
    'bar_0'
    """

    def value_factory(self, prefix: str) -> eve_utils.SequentialIDGenerator:
        return eve_utils.SequentialIDGenerator(prefix=prefix)
