# 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.
# ruff: noqa

import pytest
import numpy as np
import tilelang
from tilelang import tvm as tvm
import tvm
import tilelang.testing
from tvm import tirx
from tvm.script import tirx as T, ir as I


def _find_assignment(stmt, var_name):
    result = None

    def _visitor(node):
        nonlocal result
        if isinstance(node, tvm.tirx.Bind) and node.var.name == var_name:
            result = node

    tvm.tirx.stmt_functor.post_order_visit(stmt, _visitor)
    return result


def _find_compute_scope(func):
    result = None

    def _visitor(stmt):
        if isinstance(stmt, tirx.AttrStmt) and stmt.attr_key == "compute_scope":
            nonlocal result
            result = stmt

    tirx.stmt_functor.post_order_visit(func.body, _visitor)

    return result


def _collect_nodes(stmt, node_type):
    nodes = []

    def _visitor(node):
        if isinstance(node, node_type):
            nodes.append(node)

    tvm.tirx.stmt_functor.post_order_visit(stmt, _visitor)
    return nodes


@pytest.mark.parametrize("use_global_symbol", [False])
def test_no_op_when_global_symbol_is_absent(use_global_symbol):
    func_attr = {"target": tvm.target.Target("llvm", host="llvm")}

    @T.prim_func(private=True)
    def before():
        T.func_attr(func_attr)
        T.evaluate(0)

    if use_global_symbol:
        before = before.with_attr("global_symbol", "main")

    after = tilelang.transform.MakePackedAPI()(tvm.IRModule.from_expr(before))["main"]
    if use_global_symbol:
        assert len(after.params) == 4
    else:
        tvm.ir.assert_structural_equal(before, after)


def test_target_host_removed():
    """After MakePackedAPI, host-side target should be the host

    MakePackedAPI is the last transform that requires both the device
    and the host.  After MakePackedAPI, the target attribute should
    only contain the host-side target.
    """

    host = tvm.target.Target("llvm")

    @I.ir_module
    class before:
        @T.prim_func
        def main(A: T.Buffer(1, "float32")):
            T.func_attr({"global_symbol": "main", "target": T.target("cuda", host=host)})
            T.evaluate(0)

    after = tilelang.transform.MakePackedAPI()(before)
    target_attr = after["main"].attrs["target"]
    assert str(host) == str(target_attr)


def test_internal_subroutine_call():
    """Internal subroutines should not use the PackedFunc API

    A subroutine without the "global_symbol" attribute is an internal
    subroutine, and is not directly exposed to a user of the generated
    `runtime.Module`.  Therefore, it doesn't need to follow the
    PackedFunc API.
    """

    @I.ir_module
    class before:
        @T.prim_func
        def main(A: T.Buffer(1, "float32")):
            T.func_attr({"target": T.target("llvm", host="llvm")})
            before.subroutine(A.data)

        # this test fails if it's made public
        @T.prim_func(private=True)
        def subroutine(A_data: T.handle("float32")):
            T.func_attr({"target": T.target("llvm")})
            T.evaluate(A_data)

    after = tilelang.transform.MakePackedAPI()(before)
    tvm.ir.assert_structural_equal(before["subroutine"], after["subroutine"])

    compute_scope = _find_compute_scope(after["main"])
    subroutine_call_op = compute_scope.body.value.op
    assert isinstance(subroutine_call_op, tvm.ir.GlobalVar), (
        f"The main function's CallNode should use the subroutine's GLobalVar as the operation, "
        f"but instead has an operation of type {subroutine_call_op}"
    )


def test_assume_requires_runtime_check_is_lowered_to_assert():
    n = tirx.Var("n", "int32")
    condition = n % 4 == 0
    message = "n must be divisible by 4"
    body = tirx.AttrStmt(
        condition,
        "tl.assume",
        tirx.StringImm(message),
        tirx.Evaluate(0),
    )
    body = tirx.AttrStmt(
        condition,
        "tl.assume_requires_runtime_check",
        tirx.StringImm(message),
        body,
    )
    body = tirx.AttrStmt(
        condition,
        "tl.assume",
        tirx.StringImm("compiler-generated assumption"),
        body,
    )
    before = tirx.PrimFunc([n], body).with_attr("global_symbol", "main")
    before = before.with_attr("target", tvm.target.Target("cuda", host="llvm"))

    after_mod = tilelang.transform.MakePackedAPI()(tvm.IRModule.from_expr(before))
    after = after_mod["main"]
    assert not [stmt for stmt in _collect_nodes(after.body, tirx.AttrStmt) if stmt.attr_key == "tl.assume_requires_runtime_check"]
    assert [stmt for stmt in _collect_nodes(after.body, tirx.AttrStmt) if stmt.attr_key == "tl.assume"]

    matching_asserts = [
        stmt for stmt in _collect_nodes(after.body, tirx.AssertStmt) if any(part.value == message for part in stmt.message_parts)
    ]
    assert len(matching_asserts) == 1
    tvm.ir.assert_structural_equal(matching_asserts[0].condition, condition, map_free_vars=True)

    simplified = tilelang.transform.Simplify()(after_mod)["main"]
    matching_asserts = [
        stmt for stmt in _collect_nodes(simplified.body, tirx.AssertStmt) if any(part.value == message for part in stmt.message_parts)
    ]
    assert len(matching_asserts) == 1
    tvm.ir.assert_structural_equal(matching_asserts[0].condition, condition, map_free_vars=True)


def test_optimizer_only_assume_is_not_lowered_to_assert():
    n = tirx.Var("n", "int32")
    message = "compiler-generated shape assumption"
    body = tirx.AttrStmt(
        n > 0,
        "tl.assume",
        tirx.StringImm(message),
        tirx.Evaluate(0),
    )
    before = tirx.PrimFunc([n], body).with_attr("global_symbol", "main")
    before = before.with_attr("target", tvm.target.Target("cuda", host="llvm"))

    after = tilelang.transform.MakePackedAPI()(tvm.IRModule.from_expr(before))["main"]
    assert [stmt for stmt in _collect_nodes(after.body, tirx.AttrStmt) if stmt.attr_key == "tl.assume"]
    assert not [stmt for stmt in _collect_nodes(after.body, tirx.AssertStmt) if any(part.value == message for part in stmt.message_parts)]


def test_subroutine_call_to_externally_visible_subroutine():
    """Externally-visible subroutines should use the PackedFunc API

    Because the subroutine may be called directly by a user, it must
    use the PackedFunc API.  Its signature should be updated to the
    PackedFunc signature, and call sites should be updated to use
    `T.tvm_call_cpacked`.
    """

    @I.ir_module
    class before:
        @T.prim_func
        def main(A: T.Buffer(1, "float32")):
            T.func_attr({"global_symbol": "main", "target": T.target("llvm", host="llvm")})
            before.subroutine(A.data)

        @T.prim_func
        def subroutine(A_data: T.handle("float32")):
            T.func_attr({"global_symbol": "subroutine", "target": T.target("llvm", host="llvm")})
            T.evaluate(A_data)

    after = tilelang.transform.MakePackedAPI()(before)

    main_compute_scope = _find_compute_scope(after["main"])
    assert main_compute_scope is not None
    subroutine_compute_scope = _find_compute_scope(after["subroutine"])
    assert subroutine_compute_scope is not None

    subroutine_call_op = main_compute_scope.body.value.op
    assert isinstance(subroutine_call_op, tvm.ir.Op) and subroutine_call_op.name == "tirx.tvm_call_cpacked", (
        f"The main function's CallNode should be lowered to the builtin 'tirx.tvm_call_cpacked', "
        f"but instead has an operation of type {subroutine_call_op}"
    )


@tilelang.testing.requires_llvm
def test_function_call_with_wrong_argument_count():
    """Argument counts must be checked before accessing the type codes"""

    @T.prim_func
    def func(
        A: T.Buffer([16, 16], "int32"),
        B: T.Buffer([16, 16], "int32"),
        C: T.Buffer([16, 16], "int32"),
        D: T.Buffer([16, 16], "int32"),
    ):
        pass

    built = tvm.compile(func, target="llvm")

    with pytest.raises(tvm.TVMError):
        built()


@tilelang.testing.requires_llvm
def test_function_call_with_wrong_type_code():
    """Type codes must be checked before accessing the arguments"""

    @T.prim_func
    def func(A: T.Buffer([16, 16], "int32")):
        pass

    built = tvm.compile(func, target="llvm")

    with pytest.raises(tvm.TVMError):
        built(0)


@tilelang.testing.requires_llvm
def test_function_call_with_null_data_pointer():
    """The data pointer must be checked before accessing the array"""

    @T.prim_func
    def func(A: T.Buffer([16, 16], "int32"), B: T.Buffer([16, 16], "int32")):
        for i, j in T.grid(16, 16):
            B[i, j] = A[i, j]

    built = tvm.compile(func, target="llvm")

    A = tvm.nd.array(np.zeros([16], dtype="int32"))
    B = tvm.nd.empty([16, 16], "int32", tvm.cpu())

    with pytest.raises(tvm.TVMError):
        built(A, B)


@tilelang.testing.requires_llvm
def test_function_call_with_wrong_dimensionality():
    """The dimensionality must be checked before validating the shape"""

    @T.prim_func
    def func(A: T.Buffer([16, 16], "int32"), B: T.Buffer([16, 16], "int32")):
        for i, j in T.grid(16, 16):
            B[i, j] = A[i, j]

    built = tvm.compile(func, target="llvm")

    A = tvm.nd.array(np.zeros([16], dtype="int32"))
    B = tvm.nd.empty([16], "int32", tvm.cpu())

    with pytest.raises(tvm.TVMError):
        built(A, B)


if __name__ == "__main__":
    tilelang.testing.main()
