"""Tests for CUDA Graph Runner.

Tests capture, replay, and performance of CUDA Graph optimization
for decode phase operations.
"""

import pytest
import torch

from sagellm_backend.kernels import CUDAGraphConfig, CUDAGraphRunner


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
class TestCUDAGraphRunner:
    """Test suite for CUDA Graph Runner."""

    def simple_forward_fn(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        """Simple forward function for testing."""
        # Simulate a simple model: embedding lookup + linear
        batch_size, seq_len = input_ids.shape
        hidden_size = 128

        # Embedding
        embedding = torch.randn(50000, hidden_size, device=input_ids.device, dtype=torch.float32)
        x = torch.nn.functional.embedding(input_ids, embedding)

        # Simple linear transformation
        weight = torch.randn(hidden_size, hidden_size, device=x.device, dtype=x.dtype)
        output = torch.matmul(x, weight)

        return output

    def test_basic_capture_and_replay(self):
        """Test basic CUDA graph capture and replay."""
        config = CUDAGraphConfig(
            max_batch_size=16,
            graph_batch_sizes=[1, 2, 4, 8],
            capture_warmup_runs=1,
        )

        runner = CUDAGraphRunner(
            forward_fn=self.simple_forward_fn,
            config=config,
        )

        # Test capture for batch_size=4
        batch_size = 4
        seq_len = 1  # Decode always uses seq_len=1

        # Create sample inputs
        input_ids = torch.randint(0, 1000, (batch_size, seq_len), device="cuda")
        positions = torch.arange(seq_len, device="cuda").unsqueeze(0).expand(batch_size, -1)

        # Capture graph
        runner.capture(
            batch_size=batch_size,
            sample_inputs={
                "input_ids": input_ids,
                "positions": positions,
            },
        )

        assert batch_size in runner._graphs

        # Replay with new inputs
        new_input_ids = torch.randint(0, 1000, (batch_size, seq_len), device="cuda")
        new_positions = torch.arange(seq_len, device="cuda").unsqueeze(0).expand(batch_size, -1)

        output = runner.replay(
            batch_size=batch_size,
            inputs={
                "input_ids": new_input_ids,
                "positions": new_positions,
            },
        )

        assert output.shape == (batch_size, seq_len, 128)
        assert output.device.type == "cuda"

    def test_replay_matches_eager(self):
        """Test that graph replay produces same results as eager execution."""
        config = CUDAGraphConfig(
            max_batch_size=8,
            graph_batch_sizes=[2],
            capture_warmup_runs=1,
        )

        runner = CUDAGraphRunner(
            forward_fn=self.simple_forward_fn,
            config=config,
        )

        batch_size = 2
        seq_len = 1

        # Fix random seed for reproducibility
        torch.manual_seed(42)
        input_ids = torch.randint(0, 1000, (batch_size, seq_len), device="cuda")
        positions = torch.arange(seq_len, device="cuda").unsqueeze(0).expand(batch_size, -1)

        # Get eager output
        torch.manual_seed(42)
        eager_output = self.simple_forward_fn(input_ids, positions)

        # Capture graph
        torch.manual_seed(42)
        runner.capture(
            batch_size=batch_size,
            sample_inputs={
                "input_ids": input_ids,
                "positions": positions,
            },
        )

        # Replay graph
        torch.manual_seed(42)
        graph_output = runner.replay(
            batch_size=batch_size,
            inputs={
                "input_ids": input_ids,
                "positions": positions,
            },
        )

        # Outputs should match (within floating point tolerance)
        assert torch.allclose(eager_output, graph_output, rtol=1e-5, atol=1e-5)

    def test_multiple_batch_sizes(self):
        """Test capturing graphs for multiple batch sizes."""
        config = CUDAGraphConfig(
            max_batch_size=16,
            graph_batch_sizes=[1, 2, 4, 8],
            capture_warmup_runs=1,
        )

        runner = CUDAGraphRunner(
            forward_fn=self.simple_forward_fn,
            config=config,
        )

        # Capture for all configured batch sizes
        for batch_size in config.graph_batch_sizes:
            input_ids = torch.randint(0, 1000, (batch_size, 1), device="cuda")
            positions = torch.zeros(batch_size, 1, dtype=torch.long, device="cuda")

            runner.capture(
                batch_size=batch_size,
                sample_inputs={
                    "input_ids": input_ids,
                    "positions": positions,
                },
            )

        # Verify all graphs were captured
        assert len(runner._graphs) == len(config.graph_batch_sizes)
        for bs in config.graph_batch_sizes:
            assert bs in runner._graphs

    def test_capture_all_utility(self):
        """Test capture_all convenience method."""
        config = CUDAGraphConfig(
            max_batch_size=8,
            graph_batch_sizes=[1, 2, 4],
            capture_warmup_runs=1,
        )

        runner = CUDAGraphRunner(
            forward_fn=self.simple_forward_fn,
            config=config,
        )

        # Use capture_all
        def get_sample_inputs(batch_size: int) -> dict:
            return {
                "input_ids": torch.randint(0, 1000, (batch_size, 1), device="cuda"),
                "positions": torch.zeros(batch_size, 1, dtype=torch.long, device="cuda"),
            }

        runner.capture_all(sample_input_fn=get_sample_inputs)

        # Verify all graphs captured
        assert len(runner._graphs) == len(config.graph_batch_sizes)

    def test_fallback_to_eager_for_uncaptured_batch(self):
        """Test fallback to eager execution for uncaptured batch sizes."""
        config = CUDAGraphConfig(
            max_batch_size=16,
            graph_batch_sizes=[2, 4],  # Only capture these
            capture_warmup_runs=1,
        )

        runner = CUDAGraphRunner(
            forward_fn=self.simple_forward_fn,
            config=config,
        )

        # Capture for batch_size=2
        input_ids = torch.randint(0, 1000, (2, 1), device="cuda")
        positions = torch.zeros(2, 1, dtype=torch.long, device="cuda")

        runner.capture(
            batch_size=2,
            sample_inputs={"input_ids": input_ids, "positions": positions},
        )

        # Try to replay with uncaptured batch_size=3 (should fall back to eager)
        input_ids_3 = torch.randint(0, 1000, (3, 1), device="cuda")
        positions_3 = torch.zeros(3, 1, dtype=torch.long, device="cuda")

        # Check has_graph_for_batch_size
        assert runner.has_graph_for_batch_size(2)
        assert not runner.has_graph_for_batch_size(3)

    def test_warmup_runs(self):
        """Test that warmup runs are executed before capture."""
        config = CUDAGraphConfig(
            max_batch_size=4,
            graph_batch_sizes=[2],
            capture_warmup_runs=3,  # Multiple warmups
        )

        runner = CUDAGraphRunner(
            forward_fn=self.simple_forward_fn,
            config=config,
        )

        batch_size = 2
        input_ids = torch.randint(0, 1000, (batch_size, 1), device="cuda")
        positions = torch.zeros(batch_size, 1, dtype=torch.long, device="cuda")

        # Capture should succeed even with warmup runs
        runner.capture(
            batch_size=batch_size,
            sample_inputs={"input_ids": input_ids, "positions": positions},
        )

        assert batch_size in runner._graphs

    def test_graph_memory_efficiency(self):
        """Test that graphs reuse input/output buffers."""
        config = CUDAGraphConfig(
            max_batch_size=8,
            graph_batch_sizes=[4],
            capture_warmup_runs=1,
        )

        runner = CUDAGraphRunner(
            forward_fn=self.simple_forward_fn,
            config=config,
        )

        batch_size = 4
        input_ids = torch.randint(0, 1000, (batch_size, 1), device="cuda")
        positions = torch.zeros(batch_size, 1, dtype=torch.long, device="cuda")

        # Capture
        runner.capture(
            batch_size=batch_size,
            sample_inputs={"input_ids": input_ids, "positions": positions},
        )

        # Check that buffers are stored
        graph_info = runner._graphs[batch_size]
        assert "input_ids" in graph_info.input_buffers
        assert "positions" in graph_info.input_buffers

        # Multiple replays should reuse the same buffers
        for _ in range(3):
            new_ids = torch.randint(0, 1000, (batch_size, 1), device="cuda")
            new_pos = torch.zeros(batch_size, 1, dtype=torch.long, device="cuda")

            output = runner.replay(
                batch_size=batch_size,
                inputs={"input_ids": new_ids, "positions": new_pos},
            )

            assert output.shape == (batch_size, 1, 128)


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
class TestModelRunnerWithCUDAGraph:
    """Test ModelRunnerWithCUDAGraph wrapper."""

    def test_wrapper_initialization(self):
        """Test that wrapper can be initialized."""
        from sagellm_backend.kernels.cuda_graph import ModelRunnerWithCUDAGraph

        def dummy_forward(input_ids, **kwargs):
            return torch.randn(input_ids.shape[0], 1, 128, device=input_ids.device)

        class DummyRunner:
            def __init__(self):
                self.device = "cuda"

            def decode_forward(self, input_ids, **kwargs):
                return dummy_forward(input_ids, **kwargs)

        base_runner = DummyRunner()
        config = CUDAGraphConfig(
            max_batch_size=8,
            graph_batch_sizes=[1, 2, 4],
        )

        wrapped_runner = ModelRunnerWithCUDAGraph(
            model_runner=base_runner,
            config=config,
        )

        assert wrapped_runner.model_runner is base_runner
        assert hasattr(wrapped_runner, "_graph_runner")

    def test_wrapper_automatic_graph_selection(self):
        """Test that wrapper automatically selects graph vs eager."""
        from sagellm_backend.kernels.cuda_graph import ModelRunnerWithCUDAGraph

        def dummy_forward(input_ids, **kwargs):
            return torch.randn(input_ids.shape[0], 1, 128, device=input_ids.device)

        class DummyRunner:
            def __init__(self):
                self.device = "cuda"

            def decode_forward(self, input_ids, **kwargs):
                return dummy_forward(input_ids, **kwargs)

        base_runner = DummyRunner()
        config = CUDAGraphConfig(
            max_batch_size=8,
            graph_batch_sizes=[2, 4],
            capture_warmup_runs=1,
        )

        wrapped_runner = ModelRunnerWithCUDAGraph(
            model_runner=base_runner,
            config=config,
        )

        # Capture graphs manually via the graph runner
        if wrapped_runner._graph_runner:
            for bs in config.graph_batch_sizes:
                input_ids = torch.randint(0, 1000, (bs, 1), device="cuda")
                wrapped_runner._graph_runner.capture(
                    batch_size=bs,
                    sample_inputs={"input_ids": input_ids},
                )

        # Test with captured batch size (should use graph)
        input_ids_2 = torch.randint(0, 1000, (2, 1), device="cuda")
        output_2 = wrapped_runner.decode_forward(input_ids_2)
        assert output_2.shape == (2, 1, 128)

        # Test with uncaptured batch size (should use eager)
        input_ids_3 = torch.randint(0, 1000, (3, 1), device="cuda")
        output_3 = wrapped_runner.decode_forward(input_ids_3)
        assert output_3.shape == (3, 1, 128)


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