# 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.
import pathlib

import numpy
import pytest
import tvm_ffi
import tvm_ffi.cpp
from tvm_ffi.core import TypeSchema
from tvm_ffi.module import Module


def test_build_cpp() -> None:
    cpp_path = pathlib.Path(__file__).parent.resolve() / "test_build.cc"
    output_lib_path = tvm_ffi.cpp.build(
        name="hello",
        sources=[str(cpp_path)],
    )

    mod: Module = tvm_ffi.load_module(output_lib_path)

    metadata = mod.get_function_metadata("add_one_cpu")
    assert metadata is not None, "add_one_cpu should have metadata"
    assert "type_schema" in metadata, f"{'add_one_cpu'}: {metadata}"
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[Tensor, Tensor], None]", f"{'add_one_cpu'}: {schema}"
    doc = mod.get_function_doc("add_one_cpu")
    assert doc is None

    x = numpy.array([1, 2, 3, 4, 5], dtype=numpy.float32)
    y = numpy.empty_like(x)
    mod.add_one_cpu(x, y)
    numpy.testing.assert_equal(x + 1, y)


def test_build_inline_with_metadata() -> None:  # noqa: PLR0915
    """Test functions with various input and output types."""
    # Keep module alive until all returned objects are destroyed
    mod: Module = tvm_ffi.cpp.load_inline(
        name="test_io_types",
        cpp_sources=r"""
            // int input -> int output
            int square(int x) {
                return x * x;
            }

            // float input -> float output
            float reciprocal(float x) {
                return 1.0f / x;
            }

            // bool input -> bool output
            bool negate(bool x) {
                return !x;
            }

            // String input -> String output
            tvm::ffi::String uppercase_first(tvm::ffi::String s) {
                std::string result(s.c_str());
                if (!result.empty()) {
                    result[0] = std::toupper(result[0]);
                }
                return tvm::ffi::String(result);
            }

            // Multiple inputs: int, float -> float
            float weighted_sum(int count, float weight) {
                return static_cast<float>(count) * weight;
            }

            // Multiple inputs: String, int -> String
            tvm::ffi::String repeat_string(tvm::ffi::String s, int times) {
                std::string result;
                for (int i = 0; i < times; ++i) {
                    result += s.c_str();
                }
                return tvm::ffi::String(result);
            }

            // Mixed types: bool, int, float, String -> String
            tvm::ffi::String format_data(bool flag, int count, float value, tvm::ffi::String label) {
                std::ostringstream oss;
                oss << label.c_str() << ": flag=" << (flag ? "true" : "false")
                    << ", count=" << count << ", value=" << value;
                return tvm::ffi::String(oss.str());
            }

            // Tensor input/output
            void double_tensor(tvm::ffi::TensorView input, tvm::ffi::TensorView output) {
                TVM_FFI_ICHECK(input.ndim() == 1);
                TVM_FFI_ICHECK(output.ndim() == 1);
                TVM_FFI_ICHECK(input.size(0) == output.size(0));
                DLDataType f32_dtype{kDLFloat, 32, 1};
                TVM_FFI_ICHECK(input.dtype() == f32_dtype);
                TVM_FFI_ICHECK(output.dtype() == f32_dtype);

                for (int i = 0; i < input.size(0); ++i) {
                    static_cast<float*>(output.data_ptr())[i] =
                        static_cast<const float*>(input.data_ptr())[i] * 2.0f;
                }
            }
        """,
        functions=[
            "square",
            "reciprocal",
            "negate",
            "uppercase_first",
            "weighted_sum",
            "repeat_string",
            "format_data",
            "double_tensor",
        ],
        extra_cflags=["-DTVM_FFI_DLL_EXPORT_INCLUDE_METADATA=1"],
    )

    # Test square: int -> int
    assert mod.square(5) == 25
    metadata = mod.get_function_metadata("square")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[int], int]"

    # Test reciprocal: float -> float
    result = mod.reciprocal(2.0)
    assert abs(result - 0.5) < 0.001
    metadata = mod.get_function_metadata("reciprocal")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[float], float]"

    # Test negate: bool -> bool
    assert mod.negate(True) is False
    assert mod.negate(False) is True
    metadata = mod.get_function_metadata("negate")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[bool], bool]"

    # Test uppercase_first: String -> String
    result = mod.uppercase_first("hello")
    assert result == "Hello"
    metadata = mod.get_function_metadata("uppercase_first")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[str], str]"

    # Test weighted_sum: int, float -> float
    result = mod.weighted_sum(10, 2.5)
    assert abs(result - 25.0) < 0.001
    metadata = mod.get_function_metadata("weighted_sum")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[int, float], float]"

    # Test repeat_string: String, int -> String
    result = mod.repeat_string("ab", 3)
    assert result == "ababab"
    metadata = mod.get_function_metadata("repeat_string")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[str, int], str]"

    # Test format_data: bool, int, float, String -> String
    result = mod.format_data(True, 42, 3.14, "test")
    assert "test:" in result
    assert "flag=true" in result
    assert "count=42" in result
    assert "value=3.14" in result
    metadata = mod.get_function_metadata("format_data")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[bool, int, float, str], str]"

    # Test double_tensor: Tensor, Tensor -> None
    x = numpy.array([1.0, 2.0, 3.0], dtype=numpy.float32)
    y = numpy.empty_like(x)
    mod.double_tensor(x, y)
    numpy.testing.assert_allclose(y, x * 2.0)
    metadata = mod.get_function_metadata("double_tensor")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[Tensor, Tensor], None]"


def test_build_inline_with_docstrings() -> None:
    """Test building functions with documentation using the functions dict."""
    # Keep module alive until all returned objects are destroyed
    add_docstring = (
        "Add two integers and return the sum.\n"
        "\n"
        "Parameters\n"
        "----------\n"
        "a : int\n"
        "    First integer\n"
        "b : int\n"
        "    Second integer\n"
        "\n"
        "Returns\n"
        "-------\n"
        "result : int\n"
        "    Sum of a and b"
    )

    divide_docstring = "Divides two floats. Returns a/b."

    mod: Module = tvm_ffi.cpp.load_inline(
        name="test_docs",
        cpp_sources=r"""
            int add(int a, int b) {
                return a + b;
            }

            int subtract(int a, int b) {
                return a - b;
            }

            float divide(float a, float b) {
                TVM_FFI_ICHECK(b != 0.0f) << "Division by zero";
                return a / b;
            }
        """,
        functions={
            "add": add_docstring,
            "subtract": "",  # No documentation
            "divide": divide_docstring,
        },
        extra_cflags=["-DTVM_FFI_DLL_EXPORT_INCLUDE_METADATA=1"],
    )

    # Test add function with full documentation
    assert mod.add(10, 5) == 15
    metadata = mod.get_function_metadata("add")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[int, int], int]"

    doc = mod.get_function_doc("add")
    assert doc is not None, "add should have documentation"
    assert doc == add_docstring

    # Test subtract function without documentation
    assert mod.subtract(10, 5) == 5
    metadata = mod.get_function_metadata("subtract")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[int, int], int]"

    doc = mod.get_function_doc("subtract")
    assert doc is None, "subtract should not have documentation"

    # Test divide function with short documentation
    result = mod.divide(10.0, 2.0)
    assert abs(result - 5.0) < 0.001
    metadata = mod.get_function_metadata("divide")
    assert metadata is not None
    schema = TypeSchema.from_json_str(metadata["type_schema"])
    assert str(schema) == "Callable[[float, float], float]"

    doc = mod.get_function_doc("divide")
    assert doc is not None, "divide should have documentation"
    assert doc == divide_docstring


def test_build_without_metadata() -> None:
    """Test building without metadata export."""
    mod: Module = tvm_ffi.cpp.load_inline(
        name="test_no_meta",
        cpp_sources=r"""
            // Note: NOT defining TVM_FFI_DLL_EXPORT_INCLUDE_METADATA

            int simple_add(int a, int b) {
                return a + b;
            }
        """,
        functions=["simple_add"],
    )

    # Function should still work
    result = mod.simple_add(10, 20)
    assert result == 30

    # But metadata should not be available
    metadata = mod.get_function_metadata("simple_add")
    assert metadata is None, (
        "Metadata should not be available without TVM_FFI_DLL_EXPORT_INCLUDE_METADATA"
    )

    # Doc should also not be available
    doc = mod.get_function_doc("simple_add")
    assert doc is None


if __name__ == "__main__":
    pytest.main([__file__])