# 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 functools
import pickle
import typing
from typing import TYPE_CHECKING, ClassVar, List, Optional, Union

import gt4py.eve as eve
from gt4py.eve import Coerced, SymbolName, SymbolRef, concepts, utils as eve_utils
from gt4py.eve.traits import SymbolTableTrait, ValidatedSymbolTableTrait
from gt4py.next import common
from gt4py.next.iterator.builtins import BUILTINS
from gt4py.next.type_system import type_specifications as ts


DimensionKind = common.DimensionKind


#: A custom pickler which ignores "location" and "type" attribute of IR nodes.
_semantic_node_pickler: type[pickle.Pickler] = concepts.skipping_fields_node_pickler(
    "location", "type"
)

#: Generates an unique fingerprint for IR nodes ignoring their location and type attributes.
semantic_fingerprint: concepts.NodeFingerprinter = functools.partial(
    eve_utils.content_hash, pickler=_semantic_node_pickler
)


@eve_utils.noninstantiable
class Node(eve.Node):
    location: Optional[concepts.SourceLocation] = eve.field(default=None, repr=False, compare=False)

    # TODO(tehrengruber): include in comparison if value is not None
    type: Optional[ts.TypeSpec] = eve.field(default=None, repr=False, compare=False)

    def fingerprint(self) -> str:
        """
        Generates a unique hash string for this node that is location and type agnostic.
        """
        return semantic_fingerprint(self)

    def __str__(self) -> str:
        from gt4py.next.iterator.pretty_printer import pformat

        return pformat(self)

    def __hash__(self) -> int:
        return hash(
            (
                type(self),
                *(
                    tuple(v) if isinstance(v, list) else v
                    for (k, v) in self.iter_children_items()
                    if k not in ["location", "type"]
                ),
            )
        )


class Sym(Node):  # helper
    id: Coerced[SymbolName]


@eve_utils.noninstantiable
class Expr(Node): ...


class Literal(Expr):
    value: str
    type: ts.ScalarType


class NoneLiteral(Expr):
    _none_literal: int = 0


class InfinityLiteral(Expr):
    # TODO(tehrengruber): self referential `ClassVar` not supported in eve.
    if TYPE_CHECKING:
        POSITIVE: ClassVar[InfinityLiteral]
        NEGATIVE: ClassVar[InfinityLiteral]

    name: typing.Literal["POSITIVE", "NEGATIVE"]

    def __str__(self):
        return f"{type(self).__name__}.{self.name}"


InfinityLiteral.NEGATIVE = InfinityLiteral(name="NEGATIVE")
InfinityLiteral.POSITIVE = InfinityLiteral(name="POSITIVE")


class OffsetLiteral(Expr):
    value: Union[int, str]


class AxisLiteral(Expr):
    # TODO(havogt): Refactor to use declare Axis/Dimension at the Program level.
    # Now every use of the literal has to provide the kind, where usually we only care of the name.
    value: str
    kind: common.DimensionKind = common.DimensionKind.HORIZONTAL


class SymRef(Expr):
    id: Coerced[SymbolRef]


class Lambda(Expr, SymbolTableTrait):
    params: List[Sym]
    expr: Expr


class FunCall(Expr):
    fun: Expr  # VType[Callable]
    args: List[Expr]


class FunctionDefinition(Node, SymbolTableTrait):
    id: Coerced[SymbolName]
    params: List[Sym]
    expr: Expr


class Stmt(Node): ...


class SetAt(Stmt):  # from JAX array.at[...].set()
    expr: Expr  # only `as_fieldop(stencil)(inp0, ...)` in first refactoring
    domain: Expr
    target: Expr  # `make_tuple` or SymRef


class IfStmt(Stmt):
    cond: Expr
    true_branch: list[Stmt]
    false_branch: list[Stmt]


class Temporary(Node):
    id: Coerced[eve.SymbolName]
    domain: Optional[Expr] = None
    dtype: Optional[ts.ScalarType | ts.TupleType] = None


class Program(Node, ValidatedSymbolTableTrait):
    id: Coerced[SymbolName]
    function_definitions: List[FunctionDefinition]
    params: List[Sym]
    declarations: List[Temporary]
    body: List[Stmt]

    _NODE_SYMBOLS_: ClassVar[List[Sym]] = [
        Sym(id=name) for name in sorted(BUILTINS)
    ]  # sorted for serialization stability


# TODO(fthaler): just use hashable types in nodes (tuples instead of lists)
Sym.__hash__ = Node.__hash__  # type: ignore[method-assign]
Expr.__hash__ = Node.__hash__  # type: ignore[method-assign]
Literal.__hash__ = Node.__hash__  # type: ignore[method-assign]
NoneLiteral.__hash__ = Node.__hash__  # type: ignore[method-assign]
OffsetLiteral.__hash__ = Node.__hash__  # type: ignore[method-assign]
AxisLiteral.__hash__ = Node.__hash__  # type: ignore[method-assign]
SymRef.__hash__ = Node.__hash__  # type: ignore[method-assign]
Lambda.__hash__ = Node.__hash__  # type: ignore[method-assign]
FunCall.__hash__ = Node.__hash__  # type: ignore[method-assign]
FunctionDefinition.__hash__ = Node.__hash__  # type: ignore[method-assign]
Program.__hash__ = Node.__hash__  # type: ignore[method-assign]
SetAt.__hash__ = Node.__hash__  # type: ignore[method-assign]
IfStmt.__hash__ = Node.__hash__  # type: ignore[method-assign]
InfinityLiteral.__hash__ = Node.__hash__  # type: ignore[method-assign]
Temporary.__hash__ = Node.__hash__  # type: ignore[method-assign]
