# 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.
"""Tests for the Executable class."""

import os
import tempfile

import numpy as np

import tvm
import tvm.testing
from tvm.runtime import Executable
from tvm.script import tirx as T


@tvm.script.ir_module
class MyModule:
    @T.prim_func(s_tir=True)
    def add(
        A: T.Buffer((10,), "float32"),
        B: T.Buffer((10,), "float32"),
        C: T.Buffer((10,), "float32"),
    ):
        for i in range(10):
            C[i] = A[i] + B[i]


def test_executable_init():
    """Test initialization of Executable class."""
    lib = tvm.tirx.build(MyModule, target="llvm")
    executable = Executable(lib)

    assert executable.mod is lib
    assert executable._jitted_mod is None


def test_executable_getitem():
    """Test __getitem__ method of Executable class."""
    lib = tvm.tirx.build(MyModule, target="llvm")
    executable = Executable(lib)

    # Jit the module first
    executable.jit()

    # Test __getitem__
    add_func = executable["add"]

    # Verify the function works
    a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
    b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
    c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))

    add_func(a, b, c)

    # Check results
    tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))


def test_executable_jit_already_jitted():
    """Test jit method when module is already jitted."""
    lib = tvm.tirx.build(MyModule, target="llvm")
    executable = Executable(lib)

    # First jit call
    jitted_mod1 = executable.jit()

    # Second jit call should return the cached jitted module
    jitted_mod2 = executable.jit()
    assert jitted_mod2 is jitted_mod1

    # Test with force_recompile
    jitted_mod3 = executable.jit(force_recompile=True)
    # The module might be different after force recompilation

    # Verify both modules work correctly
    a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
    b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
    c1 = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
    c2 = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))

    jitted_mod1["add"](a, b, c1)
    jitted_mod3["add"](a, b, c2)

    tvm.testing.assert_allclose(c1.numpy(), np.array([3.0] * 10, dtype="float32"))
    tvm.testing.assert_allclose(c2.numpy(), np.array([3.0] * 10, dtype="float32"))


def test_executable_export_library():
    """Test export_library method."""
    lib = tvm.tirx.build(MyModule, target="llvm")
    executable = Executable(lib)

    # Create a temporary directory for the library
    temp_dir = tempfile.mkdtemp()
    try:
        lib_path = os.path.join(temp_dir, "test_lib.so")
        executable.export_library(lib_path)

        # Verify the library was created
        assert os.path.exists(lib_path)

        # Load the library back
        loaded_mod = tvm.runtime.load_module(lib_path)
        assert loaded_mod is not None

        # Test the loaded module
        a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
        b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
        c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))

        loaded_mod["add"](a, b, c)

        # Check results
        tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))
    finally:
        # Clean up
        if os.path.exists(temp_dir):
            import shutil

            shutil.rmtree(temp_dir)


def test_executable_export_library_with_workspace():
    """Test export_library method with workspace_dir."""
    lib = tvm.tirx.build(MyModule, target="llvm")
    executable = Executable(lib)

    # Create temporary directories
    temp_dir = tempfile.mkdtemp()
    workspace_dir = tempfile.mkdtemp()

    try:
        lib_path = os.path.join(temp_dir, "test_lib.so")
        executable.export_library(lib_path, workspace_dir=workspace_dir)

        # Verify the library was created
        assert os.path.exists(lib_path)

        # Load the library back
        loaded_mod = tvm.runtime.load_module(lib_path)
        assert loaded_mod is not None

        # Test the loaded module
        a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
        b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
        c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))

        loaded_mod["add"](a, b, c)

        # Check results
        tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))
    finally:
        # Clean up
        for directory in [temp_dir, workspace_dir]:
            if os.path.exists(directory):
                import shutil

                shutil.rmtree(directory)


def test_executable_integration():
    """Integration test for Executable with a simple TVM module."""
    # Create target and build
    target = tvm.target.Target("llvm")
    lib = tvm.tirx.build(MyModule, target=target)

    # Create an executable
    executable = Executable(lib)

    # Test jit
    jitted_mod = executable.jit()
    assert jitted_mod is not None

    # Test __getitem__
    add_func = executable["add"]
    assert add_func is not None

    # Test the function works
    a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
    b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
    c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))

    add_func(a, b, c)

    # Check results
    tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))

    # Test export_library
    temp_dir = tempfile.mkdtemp()
    try:
        lib_path = os.path.join(temp_dir, "test_lib.so")
        executable.export_library(lib_path)

        # Verify the library was created
        assert os.path.exists(lib_path)

        # Load the library back
        loaded_mod = tvm.runtime.load_module(lib_path)
        assert loaded_mod is not None

        # Test the loaded module
        loaded_add = loaded_mod["add"]
        c_loaded = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
        loaded_add(a, b, c_loaded)

        # Check results
        tvm.testing.assert_allclose(c_loaded.numpy(), np.array([3.0] * 10, dtype="float32"))

    finally:
        # Clean up
        if os.path.exists(temp_dir):
            import shutil

            shutil.rmtree(temp_dir)


def test_executable_jit_force_recompile():
    """Test jit method with force_recompile=True."""
    # Create target and build
    target = tvm.target.Target("c")
    lib = tvm.tirx.build(MyModule, target=target)

    # Create an executable
    executable = Executable(lib)

    # First jit call
    jitted_mod1 = executable.jit()

    # Second jit call without force_recompile should return the same module
    jitted_mod2 = executable.jit()
    assert jitted_mod1 is jitted_mod2

    # Third jit call with force_recompile should return a new module
    jitted_mod3 = executable.jit(force_recompile=True)
    assert jitted_mod3 is not jitted_mod1

    # Test the function works
    a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
    b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
    c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))

    jitted_mod3["add"](a, b, c)

    # Check results
    tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))


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