# 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.
# pylint: disable=missing-docstring, unused-argument

import functools
import numbers
from typing import Any

import tvm_ffi

import tvm
from tvm import relax, tirx
from tvm.ir import GlobalVar
from tvm.relax import Expr, Type
from tvm.relax.script import builder as R
from tvm.relax.script.builder.frame import BindingBlockFrame
from tvm.relax.utils import convert_to_expr
from tvm.script.ir_builder import ir as I
from tvm.script.ir_builder.base import IRBuilder
from tvm.script.parser._core import Parser, dispatch, doc

from .entry import (
    MatchCastPair,
    TypeProxy,
    _normalize_ty,
    _normalize_ty_proxy,
)

relax.Expr._dispatch_type = relax.Expr  # pylint: disable=protected-access
dispatch.register_op(relax.Expr, doc.GtE, 0)(lambda lhs, rhs: lhs >= rhs)
dispatch.register_op(relax.Expr, doc.Gt, 0)(lambda lhs, rhs: lhs > rhs)
dispatch.register_op(relax.Expr, doc.LtE, 0)(lambda lhs, rhs: lhs <= rhs)
dispatch.register_op(relax.Expr, doc.Lt, 0)(lambda lhs, rhs: lhs < rhs)


def bind_assign_value(
    self: Parser,
    node: doc.expr,
    var_name: str,
    value: Any,
    anno_ty: Type | None = None,
    emit_prim_expr: bool = False,
) -> Any:
    var_table = self.var_table.get()

    if isinstance(value, tirx.Var):
        if value.name and var_name != value.name:
            self.report_error(
                node,
                "Cannot define TIR variables with different names. The LHS of binding should "
                "has the same name provided in RHS.",
            )
        if var_name in var_table:
            prev_value = var_table[var_name]
            if not isinstance(prev_value, tirx.Var):
                self.report_error(
                    node,
                    "Cannot redefine a non-TIR-variable object to a TIR variable. Please "
                    "define the TIR variable with another name.",
                )
            if prev_value.ty != value.ty:
                self.report_error(
                    node,
                    f"Expected the same dtype for TIR vars but got {value.ty} vs {prev_value.ty}",
                )
            if not isinstance(value, type(prev_value)):
                self.report_error(
                    node,
                    f"Expected the same IR type for TIR vars "
                    f"but existing value {type(value)} is mismatched "
                    f"to previous {type(prev_value)}",
                )
            value = prev_value
        IRBuilder.name(var_name, value)
        return value

    if tvm.ir.is_prim_expr(value):
        if not emit_prim_expr:
            return value

    if isinstance(value, tuple):
        value = convert_to_expr(value)
    if isinstance(value, numbers.Number):
        value = R.const(value)

    if isinstance(value, relax.Expr):
        var = R.emit(value, anno_ty)
    elif isinstance(value, MatchCastPair):
        if anno_ty is not None and not tvm_ffi.structural_equal(anno_ty, value.ty):
            self.report_error(
                node, "Cannot specify inconsistent annotation for a match cast pair. "
            )
        var = R.emit_match_cast(value.value, value.ty)
    else:
        return value

    IRBuilder.name(var_name, var)
    return var


def is_prim_value_call(node: doc.expr) -> bool:
    return isinstance(node, doc.Call) and getattr(node.func, "attr", None) == "prim_value"


def eval_ty_proxy(self: Parser, node: doc.expr) -> TypeProxy:
    try:
        annotation = self.eval_expr(node)
        return _normalize_ty_proxy(annotation)
    except Exception as err:  # pylint: disable=broad-except
        self.report_error(node, err)
        raise


def eval_ty(self: Parser, node: doc.expr, eval_str: bool = False) -> Type:
    var_table = self.var_table.get() if eval_str else None
    try:
        ty = self.eval_expr(node)
        return _normalize_ty(ty, var_table)
    except Exception as err:  # pylint: disable=broad-except
        self.report_error(node, err)
        raise


def is_called(node: Any, func_name: str) -> bool:
    # Check if it calls into a func
    if isinstance(node, doc.Call):
        # Recursive call was found
        if isinstance(node.func, doc.Name) and node.func.id == func_name:
            return True
    elif isinstance(node, list | tuple):
        for stmt in node:
            if is_called(stmt, func_name):
                return True
    elif isinstance(node, doc.AnnAssign | doc.Assign | doc.Return | doc.Expr):
        return is_called(node.value, func_name)
    elif isinstance(node, doc.With):
        return is_called(node.body, func_name)
    elif isinstance(node, doc.If):
        smts = []
        if node.body is not None:
            smts = smts + list(node.body)
        if node.orelse is not None:
            smts = smts + list(node.orelse)
        return is_called(smts, func_name)
    return False


def is_recursive(node: doc.FunctionDef) -> bool:
    # Check if it is a recursive function
    for stmt in node.body:
        if is_called(stmt, node.name):
            return True
    return False


def collect_symbolic_var_from_prelude(
    self: Parser, node: doc.FunctionDef, symbolic_vars: dict[str, tirx.Var]
) -> dict[str, tirx.Var]:
    prelude_vars = {}
    for stmt in node.body:
        if isinstance(stmt, doc.Assign) and all(
            isinstance(target, doc.Name) and target.id in symbolic_vars for target in stmt.targets
        ):
            values = self.eval_expr(stmt.value)

            try:
                iter(values)
            except TypeError:
                values = [values]

            assert len(stmt.targets) == len(values)
            for target, value in zip(stmt.targets, values):
                name = target.id
                prelude_vars[name] = value

    return {**symbolic_vars, **prelude_vars}


def collect_symbolic_var_from_params(self: Parser, node: doc.FunctionDef) -> None:
    # Collect symbolic vars from parameters
    symbolic_vars = {}
    for arg in node.args.args:
        if arg.annotation is None:
            self.report_error(arg, "Type annotation is required for function parameters.")
        param_ty_proxy = eval_ty_proxy(self, arg.annotation)

        for var_name in param_ty_proxy.get_symbolic_vars():
            if var_name not in symbolic_vars:
                symbolic_vars[var_name] = tirx.Var(var_name, "int64")

    # Update symbolic vars based on
    symbolic_vars = collect_symbolic_var_from_prelude(self, node, symbolic_vars)

    # Define symbolic vars to the current var_table frame
    for var_name, var in symbolic_vars.items():
        self.var_table.add(var_name, var, allow_shadowing=False)


@dispatch.register(token="relax", type_name="FunctionDef")
def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
    is_inner_function = self.inside_function
    self.inside_function = True

    # reserve a var for local function
    func_val = self.var_table.get().get(node.name)
    if not func_val and is_recursive(node):
        collect_symbolic_var_from_params(self, node)
        if node.returns is None:
            ret_ty = relax.TupleType([])
        else:
            ret_ty = eval_ty(self, node.returns, eval_str=True)
        params_ty = []
        for arg in node.args.args:
            if arg.annotation is None:
                self.report_error(arg, "Type annotation is required for function parameters.")
            param_ty = eval_ty(self, arg.annotation, eval_str=True)
            params_ty.append(param_ty)
        # created a var for the local function, the same var could be used for recursive call
        local_func_var = relax.Var(node.name, relax.FuncType(params_ty, ret_ty))
        self.var_table.add(node.name, local_func_var)

    purity = find_decorator_annotation(node, "pure")
    # treat the function as private if we are inside another function
    # or if it has a privacy annotation
    privacy = is_inner_function or find_decorator_annotation(node, "private", default=False)

    with self.var_table.with_frame():
        with self.with_dispatch_token("relax"):
            with R.function(is_pure=purity, is_private=privacy):
                R.func_name(node.name)
                collect_symbolic_var_from_params(self, node)

                if node.returns is not None:
                    ann_ty = eval_ty(self, node.returns, eval_str=True)
                    R.func_ret_ty(ann_ty)

                self.visit(node.args)

                for stmt in node.body:
                    if isinstance(stmt, doc.FunctionDef):
                        if not stmt.decorator_list:
                            self.report_error(stmt, "Function must be decorated")
                        dec = self.eval_expr(stmt.decorator_list[-1])
                        # inline prim_func was found
                        if dec.dispatch_token == "tirx":
                            self.report_error(stmt, "inline prim_func is disallowed in Relax IR")

                self.visit_body(node.body)
    self.inside_function = is_inner_function


def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool:
    """
    Check the value of given annotation (argument name) in the function decorator.
    Returns the value of the annotation if present, otherwise giving the default value.
    """
    # look for the named argument in the function decorator
    for dec in node.decorator_list:
        if not isinstance(dec, doc.Call) or dec.func.attr != "function":
            continue
        for keyword in dec.keywords:
            if keyword.arg == annotation:
                return keyword.value.value
    return default


@dispatch.register(token="relax", type_name="tvm_declare_function")
def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
    with self.var_table.with_frame():
        collect_symbolic_var_from_params(self, node)

        if node.returns is None:
            # Use AnyType as unknown return type
            # NOTE: Cannot use VoidType here because the return type can be refined later.
            ret_ty = relax.AnyType()
        else:
            ret_ty = eval_ty(self, node.returns, eval_str=True)
        params = []
        for arg in node.args.args:
            if arg.annotation is None:
                self.report_error(arg, "Type annotation is required for function parameters.")
            param_ty = eval_ty(self, arg.annotation, eval_str=True)
            params.append(relax.Var(arg.arg, param_ty))

    is_pure = find_decorator_annotation(node, "pure")

    func_signature = relax.Function.create_empty(params, ret_ty, is_pure=is_pure)
    return I.decl_function(node.name, func_signature)


@dispatch.register(token="relax", type_name="pre_visit_local_function")
def pre_visit_local_function(self: Parser, node: doc.Expr) -> None:
    ir_builder = IRBuilder()
    ir_builder.__enter__()


@dispatch.register(token="relax", type_name="post_visit_local_function")
def post_visit_local_function(self: Parser, node: doc.Expr) -> None:
    ir_builder = IRBuilder.current()
    result = ir_builder.get()
    ir_builder.__exit__(None, None, None)
    # reuse var if it is reserved
    reserved_var = self.var_table.get().get(node.name)
    if reserved_var:
        var = R.emit_var_binding(relax.VarBinding(reserved_var, result))
    else:
        var = R.emit(result)
    IRBuilder.name(node.name, var)
    self.var_table.add(node.name, var, allow_shadowing=False)


@dispatch.register(token="relax", type_name="Expr")
def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
    value = self.eval_expr(node.value)
    if isinstance(value, relax.Expr):
        var = R.emit(value)
        IRBuilder.name("_", var)
        is_void_value = isinstance(var.ty, relax.TupleType) and len(var.ty.fields) == 0

        if not is_void_value:
            self.report_error(
                node,
                f"Non-void relax expressions must be bound to a variable, "
                f"but expression of type {var.ty} was used as a statement.",
            )

    elif value is not None:
        self.report_error(node, f"Unsupported Expr stmt type {value}.")


@dispatch.register(token="relax", type_name="arguments")
def visit_arguments(self: Parser, node: doc.arguments) -> None:
    arg: doc.arg
    for arg in node.args:
        if arg.annotation is None:
            self.report_error(arg, "Type annotation is required for function parameters.")
        param_ty = eval_ty(self, arg.annotation, eval_str=True)
        param = R.arg(arg.arg, param_ty)

        self.var_table.add(arg.arg, param)


@dispatch.register(token="relax", type_name="tvm_annotation")
def visit_tvm_annotation(self: Parser, node: doc.expr) -> Type:
    return eval_ty(self, node, eval_str=False)


@dispatch.register(token="relax", type_name="With")
def visit_with(self: Parser, node: doc.With) -> None:
    # Currently only `with R.dataflow()` is supported
    if len(node.items) != 1:
        self.report_error(node, "Only one item is allowed.")
    item = node.items[0]
    if item.optional_vars is not None:
        self.report_error(
            item.context_expr,
            "Relax syntax doesn't allow binding expressions in `with` to variables",
        )
    frame = self.eval_expr(item.context_expr)
    with self.var_table.with_frame():
        with frame:
            self.visit(node.body)
    if isinstance(frame, BindingBlockFrame) and frame.is_dataflow:
        output_vars = frame.output_vars
        for var in output_vars:
            self.var_table.add(var.name_hint, var, allow_shadowing=True)


@dispatch.register(token="relax", type_name="Assign")
def visit_assign(self: Parser, node: doc.Assign) -> None:
    if len(node.targets) != 1:
        self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
    lhs = node.targets[0]
    rhs = self.eval_expr(node.value)
    self.eval_assign(
        target=lhs,
        source=rhs,
        bind_value=functools.partial(
            bind_assign_value,
            emit_prim_expr=is_prim_value_call(node.value),
        ),
        allow_shadowing=True,
    )


@dispatch.register(token="relax", type_name="AnnAssign")
def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
    lhs = node.target
    rhs = self.eval_expr(node.value)
    anno_ty = self.visit_tvm_annotation(node.annotation)
    self.eval_assign(
        target=lhs,
        source=rhs,
        bind_value=functools.partial(
            bind_assign_value,
            anno_ty=anno_ty,
            emit_prim_expr=is_prim_value_call(node.value),
        ),
        allow_shadowing=True,
    )


@dispatch.register(token="relax", type_name="Return")
def visit_return(self: Parser, node: doc.Assign) -> None:
    value = self.eval_expr(node.value)
    value = convert_to_expr(value)
    R.func_ret_value(value)


@dispatch.register(token="relax", type_name="If")
def visit_if(self: Parser, node: doc.If) -> None:
    if node.orelse is None:
        raise ValueError("Else statements are required for relax dialect.")
    with R.If(self.eval_expr(node.test)) as if_frame:
        with self.var_table.with_frame():
            with R.Then():
                self.visit_body(node.body)
        with self.var_table.with_frame():
            with R.Else():
                self.visit_body(node.orelse)
    self.var_table.add(if_frame.var_name, if_frame.var, allow_shadowing=True)


@dispatch.register(token="relax", type_name="enter_token")
def enter_token(self: Parser) -> dict[str, Any]:
    def relax_call(self, *args) -> Expr:
        args = [convert_to_expr(arg) if isinstance(arg, tuple) else arg for arg in args]

        if all(isinstance(x, Expr) for x in args):
            return relax.Call(self, args)
        arg_types = [type(x) for x in args]
        raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for types {arg_types}")

    context = {"GlobalVar.__call__": GlobalVar.__call__}
    GlobalVar.__call__ = relax_call
    return context


@dispatch.register(token="relax", type_name="exit_token")
def exit_token(self: Parser, context: dict[str, Any]) -> None:
    assert "GlobalVar.__call__" in context
    GlobalVar.__call__ = context.get("GlobalVar.__call__")