# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""The core parser"""

import abc
import inspect
from collections import defaultdict
from collections.abc import Callable
from contextlib import contextmanager
from typing import Any

import numpy as np

from tvm.error import DiagnosticError
from tvm.ir import GlobalVar
from tvm.runtime import Object

from ...ir_builder import IRBuilder
from . import dispatch, doc
from .diagnostics import Diagnostics, Source
from .evaluator import eval_assign, eval_expr

DEFAULT_VISIT = {
    "Interactive",
    "Module",
    "Expression",
    "Pass",
}


def _deferred(exit_f: Callable[[], None]):
    """Created context with certain exit function.

    Parameters
    ----------
    exit_f : Callable[[], None]
        The function to call when exiting the context.

    Returns
    -------
    res : Any
        The created context.
    """

    @contextmanager
    def context():
        try:
            yield
        finally:
            exit_f()

    return context()


def _do_nothing(*args, **kwargs):  # pylint: disable=unused-argument
    pass


class ScriptMacro(abc.ABC):
    """Representation of a script macro.

    This is a callable object, intended to be called from the expression evaluator.
    The evaluator is expected to insert the current parser into the environment
    undef the name given by "parser_object_name".

    Once called, the ScriptMacro object will locate the current parser, and use it
    to parse the macro's body and produce the result.

    There were two major considerations for this design:
    1. Implementing hygienic and non-hygienic macros.
    2. Implementing macros that return values.

    Macro uses in TIR are only allowed at a statement-level, and they don't produce
    any values. Parsing of such macros could easily be done by intercepting doc.Call
    nodes in the TIR parser. If a macro is a value-producing expression, then there
    may not be a direct way to intercept calls to it if it's embedded in a complex
    expression. Because macros use function-call syntax, the evaluator will try to
    call the macro object, which this design relies on to parse and evaluate the macro.
    """

    parser_object_name = "__current_script_parser__"

    def __init__(
        self,
        source: Source,
        closure_vars: dict[str, Any],
        func: Callable,
        hygienic: bool,
    ) -> None:
        self.source = source
        self.closure_vars = closure_vars
        self.func = func
        self.hygienic = hygienic

    def __repr__(self):
        return self.source.source

    @abc.abstractmethod
    def parse_macro(self, parser: "Parser") -> Any:
        """The main macro parsing function. Different scripts may have different
        ways to parse a macro, and to return a value to the evaluator.

        Parameters
        ----------
        parser : Parser
            The parser with the appropriate frame already created and populated depending
            macro's hygiene settings,

        Returns
        -------
            The return value depends on the specifics of the particular script. It can be
            "None" or any other value or any type.
        """

    def _find_parser_def(self):
        outer_frame_infos = inspect.getouterframes(inspect.currentframe())
        for finfo in outer_frame_infos:
            parser = finfo.frame.f_globals.get(ScriptMacro.parser_object_name)
            if parser is not None:
                return parser
        raise RuntimeError(f"{ScriptMacro.parser_object_name} not available")

    def get_macro_def(self):
        ast_module = self.source.as_ast()
        for decl in ast_module.body:
            if isinstance(decl, doc.FunctionDef) and decl.name == self.func.__name__:
                return decl
        raise RuntimeError(f"cannot find macro definition for {self.func.__name__}")

    def __call__(self, *args, **kwargs):
        param_binding = inspect.signature(self.func).bind(*args, **kwargs)
        param_binding.apply_defaults()
        local_vars = param_binding.arguments
        parser = self._find_parser_def()

        with parser.with_diag_source(self.source):
            if self.hygienic:
                saved_var_table = parser.var_table
                parser.var_table = VarTable()

                with parser.var_table.with_frame():
                    for k, v in self.closure_vars.items():
                        parser.var_table.add(k, v)
                    for k, v in local_vars.items():
                        parser.var_table.add(k, v)

                    parse_result = self.parse_macro(parser)

                parser.var_table = saved_var_table

            else:
                with parser.var_table.with_frame():
                    for k, v in local_vars.items():
                        parser.var_table.add(k, v)

                    parse_result = self.parse_macro(parser)

        return parse_result


class VarTableFrame:
    """The variable table frame.
    A frame of variable table stores the variables created in one block or scope.

    Parameters
    ----------
    vars : Set[str]
        The set of variable names in the variable table frame.
    """

    vars: set[str]

    def __init__(self):
        self.vars = set()

    def add(self, var: str):
        """Add a new variable into variable table frame.

        Parameters
        ----------
        var : str
            The name of new variable.
        """
        if var in self.vars:
            raise ValueError(f"Variable {var} already defined in current scope")
        self.vars.add(var)

    def pop_all(self, fn_pop: Callable[[str], None]):
        """Pop out all variable in variable table frame.

        Parameters
        ----------
        fn_pop : Callable[[str], None]
            The methods to call when popping each variable.
        """
        for var in self.vars:
            fn_pop(var)
        self.vars.clear()


class VarTable:
    """The variable table.
    A variable table stores the all variables when parsing TVMScript.

    Parameters
    ----------
    frames : List[VarTableFrame]
        The list or stack of variable table frame.

    name2value : Dict[str, List[Any]]
        The dictionary for variable table name-based query.
    """

    frames: list[VarTableFrame]
    name2value: dict[str, list[Any]]

    def __init__(self):
        self.frames = []
        self.name2value = defaultdict(list)

    def with_frame(self):
        """Create a new variable table frame as with statement.

        Returns
        -------
        res : Any
            The context with new variable table frame.
        """

        def pop_frame():
            frame = self.frames.pop()
            frame.pop_all(lambda name: self.name2value[name].pop())

        self.frames.append(VarTableFrame())
        return _deferred(pop_frame)

    def add(self, var: str, value: Any, allow_shadowing: bool = False):
        """Add a new variable to variable table.

        Parameters
        ----------
        var : str
            The name of variable.

        value : Any
            The value of variable.

        allow_shadowing : bool
            The options of whether variable shadowing allowed for this variable.
        """
        # Skip if the key and value are equal to those in the var_table
        if self.name2value[var] and isinstance(self.name2value[var][-1], type(value)):
            old_value = self.name2value[var][-1]
            if isinstance(value, np.ndarray) and (old_value == value).all():
                return
            if isinstance(old_value, Object):
                if old_value.same_as(value):
                    return
            elif old_value == value:
                return
        if allow_shadowing and var in self.frames[-1].vars:
            # Shadowing
            self.name2value[var][-1] = value
        else:
            self.frames[-1].add(var)
            self.name2value[var].append(value)

    def get(self) -> dict[str, Any]:
        """Get a variable dictionary of latest variables.

        Returns
        -------
        res : Any
            The variable dictionary copy of latest variables.
        """
        return {key: values[-1] for key, values in self.name2value.items() if values}

    def contains_in_current_frame(self, name: str) -> bool:
        """Check whether a variable name exists in the current frame."""
        return bool(self.frames) and name in self.frames[-1].vars

    def get_at_depth(self, depth: int) -> dict[str, Any]:
        """Get variables visible at the given frame depth, using current values.

        For each variable name that appears in frames 0..depth-1, count how many
        times it was pushed (to handle shadowing), then index into name2value at
        count-1 to retrieve the latest value visible at that depth.

        Parameters
        ----------
        depth : int
            The frame depth (number of frames visible).

        Returns
        -------
        res : dict[str, Any]
            Variable dictionary of values visible at the given depth.
        """
        result: dict[str, Any] = {}
        name_count: dict[str, int] = defaultdict(int)
        for frame_idx in range(min(depth, len(self.frames))):
            for name in self.frames[frame_idx].vars:
                name_count[name] += 1
        for name, count in name_count.items():
            if self.name2value[name]:
                result[name] = self.name2value[name][count - 1]
        return result

    def exist(self, value: Any) -> bool:
        """Check if any value exists in variable table.

        Parameters
        ----------
        value : Any
            The value of variable.

        Returns
        -------
        res : bool
            The existence of the value.
        """
        return any(
            value.same_as(known_value)
            for known_value_stack in self.name2value.values()
            for known_value in known_value_stack
        )


def _dispatch_wrapper(func: dispatch.ParseMethod) -> dispatch.ParseMethod:
    def _wrapper(self: "Parser", node: doc.AST) -> None:
        try:
            return func(self, node)
        except Exception as err:  # pylint: disable=broad-except
            self.report_error(node, err)
            raise

    return _wrapper


def _dispatch(self: "Parser", type_name: str) -> dispatch.ParseMethod:
    for token in [self.dispatch_tokens[-1], "default"]:
        func = dispatch.get(token=token, type_name=type_name, default=None)
        if func is not None:
            return _dispatch_wrapper(func)
    return _dispatch_wrapper(lambda self, node: self.generic_visit(node))


class Parser(doc.NodeVisitor):
    """The TVMScript parser

    Parameters
    ----------
    diag : Diagnostics
        The diagnostics for error reporting.

    dispatch_tokens : List[str]
        The list of dispatching tokens to dispatching parsing method
        of different IRs and different doc AST structure.

    var_table : VarTable
        The variable table for parsing.

    absent_params : Dict[str, None]
        Function parameters removed by a compile-time specialization.
    """

    diag: Diagnostics
    dispatch_tokens: list[str]
    function_annotations: dict[str, dict[str, Any]] | None
    absent_params: dict[str, None]
    var_table: VarTable
    inside_function: bool  # whether we are within a function
    current_class: str | None = None  # current class being parsed
    base_py_module_context: bool = False  # whether current class inherits from BasePyModule

    def __init__(
        self,
        source: Source,
        function_annotations: dict[str, dict[str, Any]],
        absent_params: dict[str, None] | None = None,
    ) -> None:
        self.diag = Diagnostics(source)
        self.dispatch_tokens = ["default"]
        self.function_annotations = function_annotations
        self.absent_params = absent_params or {}
        self.var_table = VarTable()
        self.inside_function = False

    def parse(self, extra_vars: dict[str, Any] | None = None) -> Any:
        """The main parse method for parser.

        Parameters
        ----------
        extra_vars : Optional[Dict[str, Any]]
            The optional global value table for parsing.

        Returns
        -------
        res : Any
            The doc AST node visiting result.
        """
        if extra_vars is None:
            extra_vars = {}
        with self.var_table.with_frame():
            for k, v in extra_vars.items():
                self.var_table.add(k, v)
            node = self.diag.source.as_ast()
            self.visit(node)

    def get_dispatch_token(self, node: doc.FunctionDef) -> str:
        if not isinstance(node, doc.FunctionDef):
            self.report_error(node, "Only can get dispatch token for function.")
        if not node.decorator_list:
            self.report_error(node, "Function must be decorated")
        # TODO: only the last decorator is parsed
        decorator = self.eval_expr(node.decorator_list[-1])
        if not hasattr(decorator, "dispatch_token"):
            self.report_error(node, "The parser does not understand the decorator")
        return decorator.dispatch_token

    def with_dispatch_token(self, token: str):
        """Add a new dispatching token as with statement.

        Parameters
        ----------
        token : str
            The dispatching token.

        Returns
        -------
        res : Any
            The context with new dispatching token.
        """

        self.dispatch_tokens.append(token)
        enter_func = dispatch.get(token=token, type_name="enter_token", default=lambda *args: None)
        context = enter_func(self)

        def pop_token():
            exit_func = dispatch.get(
                token=token, type_name="exit_token", default=lambda *args: None
            )
            exit_func(self, context)
            self.dispatch_tokens.pop()

        return _deferred(pop_token)

    def set_class_context(self, class_name: str, is_base_py_module: bool = False):
        """Set the current class context for parsing.

        Parameters
        ----------
        class_name : str
            The name of the current class being parsed.
        is_base_py_module : bool
            Whether the current class inherits from BasePyModule.
        """
        self.current_class = class_name
        self.base_py_module_context = is_base_py_module

    def _get_current_class_context(self) -> str | None:
        """Get the current class context.

        Returns
        -------
        Optional[str]
            The name of the current class, or None if not in a class context.
        """
        return self.current_class

    def _is_base_py_module_context(self) -> bool:
        """Check if the current class context allows Python functions.

        Returns
        -------
        bool
            True if Python functions are allowed in the current context.
        """
        return self.base_py_module_context

    def with_diag_source(self, source: Source):
        """Add a new source as with statement.

        Parameters
        ----------
        source : Source
            The source for diagnostics.

        Returns
        -------
        res : Any
            The context with new source.
        """

        last_diag = self.diag
        self.diag = Diagnostics(source)

        def pop_source():
            self.diag = last_diag

        return _deferred(pop_source)

    @contextmanager
    def with_source_span(self, node: doc.AST):
        """Make ``node``'s source range active while constructing its IR."""
        if (
            not IRBuilder.is_in_scope()
            or getattr(node, "lineno", None) is None
            or getattr(node, "col_offset", None) is None
        ):
            yield
            return
        with IRBuilder.current().with_source_span(self.diag.source.to_span(node)):
            yield

    def annotate_current_source_span(self, value: Any) -> Any:
        """Attach the active parser span to an expression result, when applicable."""
        if isinstance(value, Object) and IRBuilder.is_in_scope():
            return IRBuilder.current()._set_current_source_span(value)  # pylint: disable=protected-access
        return value

    def eval_expr(
        self,
        node: doc.Expression | doc.expr,
        extra_vars: dict[str, Any] | None = None,
    ) -> Any:
        """Expression evaluation when parsing.

        Parameters
        ----------
        node : Union[doc.expr, doc.Expression]
            The root node of AST tree node of expression to evaluate.

        extra_vars : Optional[Dict[str, Any]]
            The optional global value table for expression evaluation.

        Returns
        -------
        res : Any
            The evaluation result.
        """
        var_values = self.var_table.get()
        if extra_vars is not None:
            for k, v in extra_vars.items():
                var_values[k] = v
        var_values[ScriptMacro.parser_object_name] = self
        return eval_expr(self, node, var_values)

    def _duplicate_lhs_check(self, target: doc.expr) -> bool | set[str]:
        """Check whether duplicate lhs exists in assignment.

        Parameters
        ----------
        target : doc.expr
            The doc AST expr node for lhs.

        Returns
        -------
        res : Union[bool, Set[str]]
            The result of true if duplicate lhs exists,
            or the set of lhs names if no duplicate lhs exists.
        """
        if isinstance(target, doc.Tuple | doc.List):
            vars: set[str] = set()  # pylint: disable=redefined-builtin
            for i in target.elts:
                res = self._duplicate_lhs_check(i)
                if isinstance(res, bool) and res:
                    return True
                assert isinstance(res, set)
                if vars & res:
                    return True
                vars = vars.union(res)
            return vars
        elif isinstance(target, doc.Name):
            return {target.id}
        elif isinstance(target, doc.Starred):
            return self._duplicate_lhs_check(target.value)
        elif isinstance(target, doc.Attribute):
            # Attribute assignment like packedB.data = ..., treated as rebinding.
            return {target.attr}
        else:
            self.report_error(target, "Invalid type in assign statement")
            raise NotImplementedError

    def eval_assign(
        self,
        target: doc.expr,
        source: Any,
        bind_value: Callable[["Parser", doc.expr, str, Any], Any],
        allow_shadowing: bool = False,
    ) -> dict[str, Any]:
        """Expression assignment evaluation when parsing.

        Parameters
        ----------
        target : doc.expr
            The root node of AST tree node of assigned expression to evaluate.

        source : Any
            The source to be assigned with evaluated expression.

        bind_value : Callable[["Parser", doc.expr, str, Any], Any]
            The value binding method when assigning the values to variables.

        allow_shadowing : bool
            The options of whether variable shadowing allowed for assignment.

        Returns
        -------
        res : Dict[str, Any]
            The dictionary of assignment result.
        """
        if self._duplicate_lhs_check(target) is True:
            self.report_error(target, "Duplicate vars assigned.")
        var_values = eval_assign(self, target, source)
        for k, v in var_values.items():
            var = bind_value(self, target, k, v)
            self.var_table.add(k, var, allow_shadowing)
        return var_values

    def report_error(self, node: doc.AST, err: Exception | str) -> None:  # pylint: disable=no-self-use
        """The error reporting when parsing.

        Parameters
        ----------
        node : doc.AST
            The doc AST node with errors.

        err: Union[Exception, str]
            The error to report.
        """

        # If the error is already being raised as a DiagnosticError,
        # re-raise it without wrapping it in a DiagnosticContext.
        if isinstance(err, DiagnosticError):
            raise err

        # Only take the last line of the error message
        if isinstance(err, RuntimeError):
            lines = list(filter(None, str(err).split("\n")))
            msg = lines[-1] if lines else (str(err) or type(err).__name__)
        elif isinstance(err, KeyError):
            msg = "KeyError: " + str(err)
        else:
            msg = str(err)

        try:
            self.diag.error(node, msg)
        except Exception as diag_err:
            # Calling self.diag.error is guaranteed to throw an
            # exception.  When shown to a user, this error should
            # reference the point of error within the provided
            # TVMScript.  However, when caught in pdb, the full
            # traceback should be available for debugging.
            if isinstance(err, Exception):
                diag_err = diag_err.with_traceback(err.__traceback__)
            raise diag_err

    def visit(self, node: doc.AST) -> None:
        """The general visiting method.

        Parameters
        ----------
        node : doc.AST
            The doc AST node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        if isinstance(node, list | tuple):
            for item in node:
                self.visit(item)
            return
        if not isinstance(node, doc.AST):
            return
        name = node.__class__.__name__.split(".")[-1]
        if name in DEFAULT_VISIT:
            func = self.generic_visit
        else:
            func = getattr(self, "visit_" + name, None)
        if func is None:
            raise NotImplementedError(f"Visitor of AST node is not implemented: {name}")
        try:
            with self.with_source_span(node):
                func(node)
        except Exception as err:  # pylint: disable=broad-except
            self.report_error(node, err)

    def visit_body(self, node: list[doc.stmt]) -> Any:
        """The general body visiting method.

        Parameters
        ----------
        node : List[doc.stmt]
            The list of statements in body.

        Returns
        -------
        res : Any
            The visiting result.
        """
        for stmt in node:
            self.visit(stmt)

    def visit_tvm_annotation(self, node: doc.expr) -> Any:
        """The general TVM annotation visiting method.

        Parameters
        ----------
        node : doc.expr
            The doc AST expr node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "tvm_annotation")(self, node)

    def visit_FunctionDef(self, node: doc.FunctionDef) -> None:  # pylint: disable=invalid-name
        """The general function definition visit method.

        Parameters
        ----------
        node : doc.FunctionDef
            The doc FunctionDef node.
        """
        token = self.get_dispatch_token(node)
        func = dispatch.get(token=token, type_name="FunctionDef", default=None)
        if func is None:
            self.report_error(
                node,
                """The parser does not understand the decorator,
                or visit_FunctionDef is not implemented for the decorator with token: """
                + token,
            )
        _dispatch(self, "pre_visit_local_function")(self, node)
        _dispatch_wrapper(func)(self, node)
        _dispatch(self, "post_visit_local_function")(self, node)

    def visit_tvm_declare_function(self, node: doc.FunctionDef) -> GlobalVar:
        token = self.get_dispatch_token(node)
        with self.with_dispatch_token(token):
            return _dispatch(self, "tvm_declare_function")(self, node)

    def visit_ClassDef(self, node: doc.ClassDef) -> Any:  # pylint: disable=invalid-name
        """The general class definition visiting method.

        Parameters
        ----------
        node : doc.ClassDef
            The doc AST class definition node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        func = dispatch.get(token="ir", type_name="ClassDef", default=None)
        if func is None:
            self.report_error(node, "The parser does not understand the decorator")
        _dispatch_wrapper(func)(self, node)

    def visit_arguments(self, node: doc.arguments) -> Any:
        """The general arguments visiting method.

        Parameters
        ----------
        node : doc.arguments
            The doc AST arguments node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "arguments")(self, node)

    def visit_For(self, node: doc.For) -> Any:  # pylint: disable=invalid-name
        """The general for visiting method.

        Parameters
        ----------
        node : doc.For
            The doc AST for node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "For")(self, node)

    def visit_While(self, node: doc.While) -> Any:  # pylint: disable=invalid-name
        """The general while visiting method.

        Parameters
        ----------
        node : doc.While
            The doc AST while node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "While")(self, node)

    def visit_With(self, node: doc.With) -> Any:  # pylint: disable=invalid-name
        """The general with visiting method.

        Parameters
        ----------
        node : doc.With
            The doc AST with node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "With")(self, node)

    def visit_Assign(self, node: doc.Assign) -> Any:  # pylint: disable=invalid-name
        """The general assign visiting method.

        Parameters
        ----------
        node : doc.Assign
            The doc AST assign node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "Assign")(self, node)

    def visit_AnnAssign(self, node: doc.AnnAssign) -> Any:  # pylint: disable=invalid-name
        """The general annotated assign visiting method.

        Parameters
        ----------
        node : doc.Assign
            The doc AST annotated assign node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "AnnAssign")(self, node)

    def visit_Expr(self, node: doc.Expr) -> Any:  # pylint: disable=invalid-name
        """The general expression visiting method.

        Parameters
        ----------
        node : doc.Expr
            The doc AST expression node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "Expr")(self, node)

    def visit_If(self, node: doc.If) -> Any:  # pylint: disable=invalid-name
        """The general if visiting method.

        Parameters
        ----------
        node : doc.If
            The doc AST if node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "If")(self, node)

    def visit_AugAssign(self, node: doc.AugAssign) -> Any:  # pylint: disable=invalid-name
        """The general augmented assignment visiting method.

        Parameters
        ----------
        node : doc.AugAssign
            The doc AST augmented assignment node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "AugAssign")(self, node)

    def visit_Assert(self, node: doc.Assert) -> Any:  # pylint: disable=invalid-name
        """The general assert visiting method.

        Parameters
        ----------
        node : doc.Assert
            The doc AST assert node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "Assert")(self, node)

    def visit_Return(self, node: doc.Return) -> Any:  # pylint: disable=invalid-name
        """The general return visiting method.

        Parameters
        ----------
        node : doc.Return
            The doc AST return node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "Return")(self, node)

    def visit_Continue(self, node: doc.Continue) -> Any:  # pylint: disable=invalid-name
        """The general continue visiting method.

        Parameters
        ----------
        node : doc.Continue
            The doc AST continue node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "Continue")(self, node)

    def visit_Break(self, node: doc.Break) -> Any:  # pylint: disable=invalid-name
        """The general break visiting method.

        Parameters
        ----------
        node : doc.Break
            The doc AST break node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "Break")(self, node)

    def visit_Nonlocal(self, node: doc.Nonlocal) -> Any:  # pylint: disable=invalid-name
        """The general nonlocal visiting method.

        Parameters
        ----------
        node : doc.Nonlocal
            The doc AST nonlocal node.

        Returns
        -------
        res : Any
            The visiting result.
        """
        return _dispatch(self, "Nonlocal")(self, node)
