"""Tests for CUDA hardware-specific optimization pass."""

from __future__ import annotations

from sagellm_backend.graph.graph import Graph, Node, OpType
from sagellm_backend.graph.passes.cuda_hardware_optimization import CUDAHardwareOptimizationPass


class TestCUDAHardwareOptimizationPass:
    """Test suite for CUDA hardware-specific optimization pass."""

    def test_pass_name(self) -> None:
        """Pass should expose a stable name."""
        pass_obj = CUDAHardwareOptimizationPass()
        assert pass_obj.name == "cuda_hardware_optimization"

    def test_skip_non_cuda_graph(self) -> None:
        """CPU graph should not be modified by CUDA-specific pass."""
        pass_obj = CUDAHardwareOptimizationPass()
        graph = Graph(metadata={"target_device": "cpu"})

        x = Node(node_id="x", op_type=OpType.INPUT, device="cpu")
        y = Node(node_id="y", op_type=OpType.OUTPUT, device="cpu")
        graph.add_node(x)
        graph.add_node(y)
        graph.add_edge("x", "y")
        graph.outputs = ["y"]

        result = pass_obj.run(graph)

        assert result.modified is False
        assert "cuda_autotune" not in graph.metadata

    def test_flash_attention_annotation(self) -> None:
        """CUDA attention node should be annotated for FlashAttention fusion."""
        pass_obj = CUDAHardwareOptimizationPass()
        graph = Graph(metadata={"target_device": "cuda"})

        qkv = Node(node_id="qkv", op_type=OpType.INPUT, device="cuda:0", dtype="float16")
        attn = Node(
            node_id="attn",
            op_type=OpType.ATTENTION,
            device="cuda:0",
            dtype="float16",
            attrs={"head_dim": 128},
        )
        out = Node(node_id="out", op_type=OpType.OUTPUT, device="cuda:0", dtype="float16")

        for node in [qkv, attn, out]:
            graph.add_node(node)

        graph.add_edge("qkv", "attn")
        graph.add_edge("attn", "out")
        graph.outputs = ["out"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert graph.nodes["attn"].attrs["cuda_fusion"] == "flash_attention"
        assert graph.nodes["attn"].metadata["kernel_impl"] == "flash_attention"
        assert result.metrics["fused_flash_attention"] >= 1

    def test_fuse_gemm_bias_relu(self) -> None:
        """GEMM + Bias + ReLU chain should be fused into one fusion node."""
        pass_obj = CUDAHardwareOptimizationPass()
        graph = Graph(metadata={"target_device": "cuda"})

        x = Node(node_id="x", op_type=OpType.INPUT, device="cuda:0", dtype="float16")
        w = Node(node_id="w", op_type=OpType.INPUT, device="cuda:0", dtype="float16")
        bias = Node(node_id="bias", op_type=OpType.CONSTANT, device="cuda:0", dtype="float16")
        gemm = Node(
            node_id="gemm",
            op_type=OpType.MATMUL,
            device="cuda:0",
            dtype="float16",
            attrs={"shape": [32, 4096, 4096]},
        )
        add = Node(node_id="add", op_type=OpType.ADD, device="cuda:0", dtype="float16")
        relu = Node(node_id="relu", op_type=OpType.RELU, device="cuda:0", dtype="float16")
        out = Node(node_id="out", op_type=OpType.OUTPUT, device="cuda:0", dtype="float16")

        for node in [x, w, bias, gemm, add, relu, out]:
            graph.add_node(node)

        graph.add_edge("x", "gemm")
        graph.add_edge("w", "gemm")
        graph.add_edge("gemm", "add")
        graph.add_edge("bias", "add")
        graph.add_edge("add", "relu")
        graph.add_edge("relu", "out")
        graph.outputs = ["out"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert "gemm" not in graph.nodes
        assert "add" not in graph.nodes
        assert graph.nodes["relu"].op_type == OpType.FUSION
        assert graph.nodes["relu"].attrs["fused_pattern"] == "gemm_bias_relu"
        assert result.metrics["fused_gemm_bias_relu"] >= 1

    def test_fuse_layernorm_dropout(self) -> None:
        """LayerNorm + Dropout (represented via MUL) should be fused."""
        pass_obj = CUDAHardwareOptimizationPass()
        graph = Graph(metadata={"target_device": "cuda"})

        x = Node(node_id="x", op_type=OpType.INPUT, device="cuda:0", dtype="float16")
        ln = Node(node_id="ln", op_type=OpType.LAYER_NORM, device="cuda:0", dtype="float16")
        dropout = Node(
            node_id="dropout",
            op_type=OpType.MUL,
            device="cuda:0",
            dtype="float16",
            attrs={"dropout_prob": 0.1},
        )
        out = Node(node_id="out", op_type=OpType.OUTPUT, device="cuda:0", dtype="float16")

        for node in [x, ln, dropout, out]:
            graph.add_node(node)

        graph.add_edge("x", "ln")
        graph.add_edge("ln", "dropout")
        graph.add_edge("dropout", "out")
        graph.outputs = ["out"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert "ln" not in graph.nodes
        assert graph.nodes["dropout"].op_type == OpType.FUSION
        assert graph.nodes["dropout"].attrs["fused_pattern"] == "layernorm_dropout"
        assert result.metrics["fused_layernorm_dropout"] >= 1

    def test_kernel_autotune_select_and_cache(self) -> None:
        """Auto-tuning should select best config and reuse cache on repeated shapes."""
        pass_obj = CUDAHardwareOptimizationPass()
        graph = Graph(metadata={"target_device": "cuda"})

        x = Node(node_id="x", op_type=OpType.INPUT, device="cuda:0", dtype="float16")
        w = Node(node_id="w", op_type=OpType.INPUT, device="cuda:0", dtype="float16")
        gemm1 = Node(
            node_id="gemm1",
            op_type=OpType.MATMUL,
            device="cuda:0",
            dtype="float16",
            attrs={
                "shape": [16, 1024, 1024],
                "autotune_candidates": ["tile_64x64", "tile_128x128"],
                "autotune_profile_ms": {"tile_64x64": 1.7, "tile_128x128": 1.2},
            },
        )
        gemm2 = Node(
            node_id="gemm2",
            op_type=OpType.MATMUL,
            device="cuda:0",
            dtype="float16",
            attrs={
                "shape": [16, 1024, 1024],
                "autotune_candidates": ["tile_64x64", "tile_128x128"],
            },
        )
        out = Node(node_id="out", op_type=OpType.OUTPUT, device="cuda:0", dtype="float16")

        for node in [x, w, gemm1, gemm2, out]:
            graph.add_node(node)

        graph.add_edge("x", "gemm1")
        graph.add_edge("w", "gemm1")
        graph.add_edge("x", "gemm2")
        graph.add_edge("w", "gemm2")
        graph.add_edge("gemm2", "out")
        graph.outputs = ["out"]

        result = pass_obj.run(graph)

        cache = graph.metadata["cuda_kernel_autotune_cache"]
        assert result.modified is True
        assert len(cache) >= 1
        assert gemm1.metadata["autotune_config"] == "tile_128x128"
        assert gemm2.metadata["autotune_cache_hit"] is True

    def test_tensorrt_and_cuda_graph_metadata(self) -> None:
        """Pass should prepare TensorRT and CUDA Graph metadata for runtime integration."""
        pass_obj = CUDAHardwareOptimizationPass()
        graph = Graph(
            metadata={
                "target_device": "cuda",
                "decode_static_shapes": True,
                "tensorrt_engine_path": "/tmp/engine.plan",
                "calibration_dataset": "calib-v1",
            }
        )

        x = Node(node_id="x", op_type=OpType.INPUT, device="cuda:0", dtype="float16")
        attn = Node(node_id="attn", op_type=OpType.ATTENTION, device="cuda:0", dtype="float16")
        out = Node(node_id="out", op_type=OpType.OUTPUT, device="cuda:0", dtype="float16")

        for node in [x, attn, out]:
            graph.add_node(node)

        graph.add_edge("x", "attn")
        graph.add_edge("attn", "out")
        graph.outputs = ["out"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert graph.metadata["tensorrt"]["onnx_export"]["enabled"] is True
        assert graph.metadata["tensorrt"]["engine"]["path"] == "/tmp/engine.plan"
        assert graph.metadata["tensorrt"]["engine"]["load_ready"] is True
        assert graph.metadata["cuda_graph"]["capture_enabled"] is True
        assert graph.nodes["attn"].metadata["tensorrt_compatible"] is True
        assert graph.nodes["attn"].metadata["cuda_graph_capture_safe"] is True
