"""
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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.
"""

# NOTE: Coverage supplement test — uses mock to reach compilation pipeline
# branches that require Triton/CUDA toolchain not available in unit tests.

import unittest
from unittest.mock import MagicMock, patch

import paddle
import triton.language as tl

TRITON_UTILS_V2_PATH = "fastdeploy.model_executor.ops.triton_ops.triton_utils_v2"
import fastdeploy.model_executor.ops.triton_ops.triton_utils_v2 as tu2


class TestGetValueHint(unittest.TestCase):
    """Test the helper function get_value_hint from triton_utils_v2."""

    def test_get_value_hint_int_and_float(self):
        """Ensure get_value_hint handles mixed int and float values."""
        vals = [10, 1, -3, 1.5]
        hint = tu2.get_value_hint(vals)
        self.assertEqual(hint, "i64,i64,i64,fp32,")


class TestKernelInterfaceV2(unittest.TestCase):
    """Test cases for KernelInterface and decorator behavior."""

    def mock_kernel(self, a, b, N: tl.constexpr, K: tl.constexpr):
        return

    def test_kernel_interface_constexpr_detection(self):
        """Verify constexpr argument detection and exclusion list generation."""
        kernel_interface = tu2.KernelInterface(self.mock_kernel, other_config={})
        self.assertEqual(kernel_interface.arg_names, ["a", "b", "N", "K"])
        self.assertEqual(kernel_interface.constexprs, [2, 3])
        self.assertEqual(kernel_interface.arg_exclude_constexpr, ["a", "b"])

    @patch("paddle.distributed.get_rank", return_value=0)
    def test_decorator_cache_hit(self, _mock_rank):
        """Ensure cached compiled ops are reused without recompilation."""
        kernel_interface = tu2.KernelInterface(self.mock_kernel, other_config={})
        kernel_interface.grid = [1, 1, 1]
        op_name = "haha_N8_K16"
        cached_fn = MagicMock()
        kernel_interface.func_map[op_name] = cached_fn
        kernel_interface.decorator(1, 2, N=8, K=16)
        cached_fn.assert_called_once_with(1, 2)

    @patch("os.makedirs")
    @patch("os.getenv", return_value="/tmp/triton_cache/rank0")
    @patch("builtins.open", new_callable=MagicMock)
    @patch("importlib.import_module")
    @patch("paddle.distributed.get_rank", return_value=0)
    @patch(f"{TRITON_UTILS_V2_PATH}.build_package")
    @patch(f"{TRITON_UTILS_V2_PATH}.rename_c_to_cu")
    @patch(f"{TRITON_UTILS_V2_PATH}.multi_process_do")
    @patch(f"{TRITON_UTILS_V2_PATH}.extract_triton_kernel")
    @patch(f"{TRITON_UTILS_V2_PATH}.find_so_path")
    def test_decorator_compile_and_import(
        self,
        mock_find_so_path,
        mock_extract,
        mock_mp_do,
        mock_rename,
        mock_build,
        mock_rank,
        mock_import,
        mock_open,
        mock_getenv,
        mock_makedirs,
    ):
        original_system = tu2.os.system
        try:
            mock_system = MagicMock(return_value=0)
            tu2.os.system = mock_system
            """Test full compilation → linking → building → importing pipeline when no cached .so exists."""
            mock_find_so_path.side_effect = [
                None,
                "/tmp/triton_cache/rank0/haha_N8_K16/mock_lib.so",
            ]
            mock_module = MagicMock()
            mock_pybind_func = MagicMock()
            mock_module.haha_N8_K16_func = mock_pybind_func
            mock_import.return_value = mock_module
            mock_system.return_value = 0
            mock_mp_do.return_value = None
            mock_build.return_value = None
            mock_extract.return_value = None
            kernel_interface = tu2.KernelInterface(self.mock_kernel, other_config={})
            kernel_interface.grid = ["N * M * N"]
            kernel_interface.decorator(1, 2, N=8, K=16)
            mock_extract.assert_called_once()
            mock_mp_do.assert_called_once()
            mock_build.assert_called_once()
            mock_import.assert_any_call("haha_N8_K16_package")
            mock_pybind_func.assert_called_once_with(1, 2)
        finally:
            tu2.os.system = original_system

    @patch("os.system")
    @patch("os.makedirs")
    @patch("os.getenv", return_value="/tmp/triton_cache/rank0")
    @patch("builtins.open", new_callable=MagicMock)
    @patch("importlib.import_module")
    @patch("paddle.distributed.get_rank", return_value=0)
    @patch(f"{TRITON_UTILS_V2_PATH}.build_package")
    @patch(f"{TRITON_UTILS_V2_PATH}.rename_c_to_cu")
    @patch(f"{TRITON_UTILS_V2_PATH}.multi_process_do")
    @patch(f"{TRITON_UTILS_V2_PATH}.extract_triton_kernel")
    @patch(f"{TRITON_UTILS_V2_PATH}.find_so_path")
    @patch(f"{TRITON_UTILS_V2_PATH}.get_pointer_hint")
    def test_tensor_and_none_branch(
        self,
        mock_get_pointer_hint,
        mock_find_so_path,
        mock_extract,
        mock_mp_do,
        mock_rename,
        mock_build,
        mock_rank,
        mock_import,
        mock_open,
        mock_getenv,
        mock_makedirs,
        mock_system,
    ):
        """Ensure decorator correctly handles Tensor and None arguments during dtype and pointer analysis."""
        ki = tu2.KernelInterface(self.mock_kernel, other_config={})
        mock_find_so_path.return_value = "/tmp/triton_cache/rank0/haha_N8_K16/mock_lib.so"
        mock_module = MagicMock()
        mock_pybind_func = MagicMock()
        mock_module.haha_N8_K16_func = mock_pybind_func
        mock_import.return_value = mock_module
        ki.grid = [1, 1, 1]
        a = paddle.to_tensor([1], dtype="float32")
        b = None
        mock_get_pointer_hint.return_value = "addr_hint"
        ki.decorator(a, b, N=8, K=16)
        mock_get_pointer_hint.assert_called_once()
        dtypes_arg = mock_get_pointer_hint.call_args[0][0]
        self.assertEqual(len(dtypes_arg), 2)
        self.assertEqual(dtypes_arg[0], a.dtype)
        self.assertEqual(dtypes_arg[1], paddle.int8)
        mock_import.assert_any_call("haha_N8_K16_package")
        mock_pybind_func.assert_called_once_with(a, b)

    def test_getitem_sets_grid_and_returns_decorator(self):
        """Ensure __getitem__ sets internal grid and returns a callable decorator."""
        kernel_interface = tu2.KernelInterface(self.mock_kernel, other_config={})
        dec = kernel_interface[["unused"]]
        self.assertTrue(isinstance(kernel_interface.grid, tuple))
        self.assertIn("max_possible_num_post_padded", kernel_interface.grid[0])
        self.assertTrue(callable(dec))


class TestPaddleUseTritonV2(unittest.TestCase):
    """Tests for paddle_use_triton_v2 decorator wrapper."""

    def test_paddle_use_triton_v2_wraps_function(self):
        """Verify paddle_use_triton_v2 returns a KernelInterface with correct key arguments."""

        @tu2.paddle_use_triton_v2(other_config={"foo": 1}, key=["N", "K"])
        def my_kernel(a, N, K):
            return

        self.assertIsInstance(my_kernel, tu2.KernelInterface)
        self.assertEqual(my_kernel.key_args, ["N", "K"])


class TestKernelInterfaceUnsupportedTypes(unittest.TestCase):
    """Test assert False paths for unsupported types in KernelInterface decorator (L192, L200)."""

    def test_unsupported_constexpr_type_raises_assertion(self):
        """Passing unsupported constexpr type triggers assert with message (L192)."""

        def kernel(a, N: tl.constexpr, K: tl.constexpr):
            return

        ki = tu2.KernelInterface(kernel, other_config={})
        ki.grid = [1, 1, 1]

        # Pass a string (unsupported type) as constexpr arg 'N'
        a = paddle.to_tensor([1], dtype="float32")
        with self.assertRaises(AssertionError) as ctx:
            ki.decorator(a, N="bad_string", K=8)
        self.assertIn("Unsupported constexpr type", str(ctx.exception))
        self.assertIn("N", str(ctx.exception))

    def test_unsupported_non_constexpr_arg_type_raises_assertion(self):
        """Passing unsupported non-constexpr arg type triggers assert with message (L200)."""

        def kernel(a, b, N: tl.constexpr, K: tl.constexpr):
            return

        ki = tu2.KernelInterface(kernel, other_config={})
        ki.grid = [1, 1, 1]

        # 'a' is a Tensor, 'b' is a non-constexpr non-Tensor — pass a string
        a = paddle.to_tensor([1], dtype="float32")
        with self.assertRaises(AssertionError) as ctx:
            ki.decorator(a, "bad_string", N=8, K=16)
        self.assertIn("Unsupported arg type", str(ctx.exception))
        self.assertIn("b", str(ctx.exception))


if __name__ == "__main__":
    unittest.main()
