"""Tests for ComputeCommOverlapStrategy."""

from __future__ import annotations

import pytest

from sagellm_core.runtime_optimizer import (
    ComputeCommOverlapStrategy,
    RuntimeOptimizer,
    RuntimeStats,
)


class MockBatch:
    """Mock batch item for testing."""

    def __init__(
        self,
        batch_type: str = "mixed",
        num_requests: int = 1,
        op_type: str | None = None,
        metadata: dict | None = None,
    ) -> None:
        self.batch_type = batch_type
        self.num_requests = num_requests
        self.op_type = op_type
        self.metadata = metadata or {}

    def is_prefill(self) -> bool:
        return self.batch_type == "prefill"

    def is_decode(self) -> bool:
        return self.batch_type == "decode"


class MockBackend:
    """Mock backend for testing."""

    def __init__(self) -> None:
        self.comm_backend = None
        self.topology = None
        self.parallelism_strategy = {}

    def set_parallelism_strategy(self, tp_degree: int, pp_degree: int) -> None:
        self.parallelism_strategy = {"tp_degree": tp_degree, "pp_degree": pp_degree}

    def get_comm_backend(self) -> None:
        return self.comm_backend

    def get_topology(self) -> MockTopology | None:
        return self.topology


class MockTopology:
    """Mock topology for testing."""

    def __init__(self, interconnect: str = "nvlink") -> None:
        self.node_interconnect = interconnect


# Test 1: Strategy initialization with default parameters
def test_compute_comm_overlap_strategy_init_default() -> None:
    """Test strategy initialization with default parameters."""
    strategy = ComputeCommOverlapStrategy()

    assert strategy.name == "compute-comm-overlap"
    assert strategy.enabled is True
    assert strategy.enable_overlap_scheduling is True
    assert strategy.enable_dynamic_tp_pp is True
    assert strategy.enable_pipeline_tuning is True
    assert strategy.tp_degree == 1
    assert strategy.pp_degree == 1
    assert strategy.base_microbatch_size == 4
    assert strategy.overlap_threshold == 0.25
    assert strategy.tp_pp_rebalance_interval == 100


# Test 2: Strategy initialization with custom parameters
def test_compute_comm_overlap_strategy_init_custom() -> None:
    """Test strategy initialization with custom parameters."""
    strategy = ComputeCommOverlapStrategy(
        enable_overlap_scheduling=False,
        enable_dynamic_tp_pp=False,
        enable_pipeline_tuning=False,
        tp_degree=4,
        pp_degree=2,
        base_microbatch_size=8,
        overlap_threshold=0.3,
        tp_pp_rebalance_interval=50,
        name="custom-overlap",
        enabled=False,
    )

    assert strategy.name == "custom-overlap"
    assert strategy.enabled is False
    assert strategy.enable_overlap_scheduling is False
    assert strategy.enable_dynamic_tp_pp is False
    assert strategy.enable_pipeline_tuning is False
    assert strategy.tp_degree == 4
    assert strategy.pp_degree == 2
    assert strategy.base_microbatch_size == 8
    assert strategy.overlap_threshold == 0.3
    assert strategy.tp_pp_rebalance_interval == 50


# Test 3: Parameter validation
def test_compute_comm_overlap_strategy_validation() -> None:
    """Test parameter validation."""
    with pytest.raises(ValueError, match="tp_degree must be >= 1"):
        ComputeCommOverlapStrategy(tp_degree=0)

    with pytest.raises(ValueError, match="pp_degree must be >= 1"):
        ComputeCommOverlapStrategy(pp_degree=0)

    with pytest.raises(ValueError, match="base_microbatch_size must be >= 1"):
        ComputeCommOverlapStrategy(base_microbatch_size=0)

    with pytest.raises(ValueError, match="overlap_threshold must be in"):
        ComputeCommOverlapStrategy(overlap_threshold=1.5)

    with pytest.raises(ValueError, match="tp_pp_rebalance_interval must be >= 1"):
        ComputeCommOverlapStrategy(tp_pp_rebalance_interval=0)


# Test 4: Overlap scheduling optimization
def test_overlap_scheduling_reordering() -> None:
    """Test overlap scheduling reorders operations correctly."""
    strategy = ComputeCommOverlapStrategy(
        enable_overlap_scheduling=True,
        enable_dynamic_tp_pp=False,
        enable_pipeline_tuning=False,
    )
    backend = MockBackend()
    strategy.set_backend(backend)
    stats = RuntimeStats()

    # Create mixed batch with different communication intensities
    step = [
        MockBatch(op_type="all_reduce"),  # comm_intensity = 1.0 (comm-heavy)
        MockBatch(op_type="attention"),  # comm_intensity = 0.1 (mixed, not 0.0)
        MockBatch(batch_type="prefill"),  # comm_intensity = 0.5 (mixed)
        MockBatch(op_type="matmul"),  # comm_intensity = 0.1 (mixed, not 0.0)
    ]

    optimized = strategy.optimize(step, stats)

    # Check reordering: mixed (low comm) -> mixed (higher comm) -> comm-heavy
    assert len(optimized) == 4

    # Check that comm-heavy op comes last
    assert optimized[-1].op_type == "all_reduce"

    # Check that compute ops come before comm-heavy
    compute_indices = [
        i for i, item in enumerate(optimized) if item.op_type in ("attention", "matmul")
    ]
    all_reduce_index = next(i for i, item in enumerate(optimized) if item.op_type == "all_reduce")

    # All compute ops should come before all_reduce
    assert all(idx < all_reduce_index for idx in compute_indices)

    # Check metadata
    for item in optimized:
        assert item.metadata.get("overlap_scheduling_applied") is True


# Test 5: Dynamic TP/PP adjustment with NVLink
def test_dynamic_tp_pp_adjustment_nvlink() -> None:
    """Test TP/PP adjustment favors TP on NVLink topology."""
    strategy = ComputeCommOverlapStrategy(
        enable_overlap_scheduling=False,
        enable_dynamic_tp_pp=True,
        enable_pipeline_tuning=False,
        tp_degree=1,
        pp_degree=2,
        overlap_threshold=0.2,
        tp_pp_rebalance_interval=10,
    )
    backend = MockBackend()
    backend.topology = MockTopology(interconnect="nvlink")
    strategy.set_backend(backend)
    stats = RuntimeStats()

    # Simulate high communication overhead
    for _ in range(30):
        strategy.observe_execution(
            batch_size=8,
            latency_ms=100.0,
            metadata={"comm_overhead_ratio": 0.3},  # Above threshold
        )

    # Trigger rebalancing by running optimize for interval steps
    step = [MockBatch()]
    for i in range(11):
        strategy.optimize(step, stats)

    # TP should increase (NVLink favors TP)
    assert strategy.tp_degree > 1
    assert backend.parallelism_strategy["tp_degree"] > 1


# Test 6: Dynamic TP/PP adjustment with InfiniBand
def test_dynamic_tp_pp_adjustment_infiniband() -> None:
    """Test TP/PP adjustment favors PP on InfiniBand topology."""
    strategy = ComputeCommOverlapStrategy(
        enable_overlap_scheduling=False,
        enable_dynamic_tp_pp=True,
        enable_pipeline_tuning=False,
        tp_degree=2,
        pp_degree=1,
        overlap_threshold=0.2,
        tp_pp_rebalance_interval=10,
    )
    backend = MockBackend()
    backend.topology = MockTopology(interconnect="ib")
    strategy.set_backend(backend)
    stats = RuntimeStats()

    # Simulate high communication overhead
    for _ in range(30):
        strategy.observe_execution(
            batch_size=8,
            latency_ms=100.0,
            metadata={"comm_overhead_ratio": 0.3},
        )

    # Trigger rebalancing
    step = [MockBatch()]
    for i in range(11):
        strategy.optimize(step, stats)

    # PP should increase (IB favors PP)
    assert strategy.pp_degree > 1
    assert backend.parallelism_strategy["pp_degree"] > 1


# Test 7: Pipeline microbatch tuning - high bubble time
def test_pipeline_tuning_high_bubble_time() -> None:
    """Test microbatch size increases when bubble time is high."""
    strategy = ComputeCommOverlapStrategy(
        enable_overlap_scheduling=False,
        enable_dynamic_tp_pp=False,
        enable_pipeline_tuning=True,
        pp_degree=4,  # Enable PP
        base_microbatch_size=4,
    )
    backend = MockBackend()
    strategy.set_backend(backend)
    stats = RuntimeStats()

    # Simulate high bubble time
    for _ in range(30):
        strategy.observe_execution(
            batch_size=8,
            latency_ms=100.0,
            metadata={"bubble_time_ratio": 0.2},  # High bubble time
        )

    step = [MockBatch()]
    optimized = strategy.optimize(step, stats)

    # Microbatch size should increase
    assert strategy.base_microbatch_size > 4
    assert optimized[0].metadata.get("pipeline_microbatch_size") == strategy.base_microbatch_size


# Test 8: Pipeline microbatch tuning - high comm overhead
def test_pipeline_tuning_high_comm_overhead() -> None:
    """Test microbatch size decreases when comm overhead is high."""
    strategy = ComputeCommOverlapStrategy(
        enable_overlap_scheduling=False,
        enable_dynamic_tp_pp=False,
        enable_pipeline_tuning=True,
        pp_degree=4,
        base_microbatch_size=8,
        overlap_threshold=0.2,
    )
    backend = MockBackend()
    strategy.set_backend(backend)
    stats = RuntimeStats()

    # Simulate high comm overhead
    for _ in range(30):
        strategy.observe_execution(
            batch_size=8,
            latency_ms=100.0,
            metadata={"comm_overhead_ratio": 0.3},  # High comm overhead
        )

    step = [MockBatch()]
    _ = strategy.optimize(step, stats)

    # Microbatch size should decrease
    assert strategy.base_microbatch_size < 8


# Test 9: Pipeline tuning disabled when PP=1
def test_pipeline_tuning_disabled_no_pp() -> None:
    """Test pipeline tuning is skipped when PP degree is 1."""
    strategy = ComputeCommOverlapStrategy(
        enable_pipeline_tuning=True,
        pp_degree=1,  # No pipeline parallelism
        base_microbatch_size=4,
    )
    backend = MockBackend()
    strategy.set_backend(backend)
    stats = RuntimeStats()

    step = [MockBatch()]
    optimized = strategy.optimize(step, stats)

    # Microbatch size should not be added to metadata when PP=1
    assert "pipeline_microbatch_size" not in optimized[0].metadata


# Test 10: Metrics tracking
def test_metrics_tracking() -> None:
    """Test metrics are tracked correctly."""
    strategy = ComputeCommOverlapStrategy()
    backend = MockBackend()
    strategy.set_backend(backend)

    # Simulate execution with metrics
    for i in range(50):
        strategy.observe_execution(
            batch_size=8,
            latency_ms=100.0,
            metadata={
                "comm_overhead_ratio": 0.2 + i * 0.001,
                "bubble_time_ratio": 0.1,
                "overlap_efficiency": 0.8,
            },
        )

    metrics = strategy.get_metrics()

    assert metrics["step_count"] == 0  # No optimize() calls yet
    assert metrics["comm_overhead_avg"] > 0.0
    assert metrics["throughput_avg"] > 0.0
    assert metrics["bubble_time_avg"] > 0.0
    assert metrics["overlap_efficiency_avg"] > 0.0
    assert "throughput_improvement" in metrics
    assert "comm_overhead_reduction" in metrics


# Test 11: Combined optimization
def test_combined_optimization() -> None:
    """Test all optimizations work together."""
    strategy = ComputeCommOverlapStrategy(
        enable_overlap_scheduling=True,
        enable_dynamic_tp_pp=True,
        enable_pipeline_tuning=True,
        pp_degree=2,
        tp_pp_rebalance_interval=5,
    )
    backend = MockBackend()
    backend.topology = MockTopology(interconnect="nvlink")
    strategy.set_backend(backend)
    stats = RuntimeStats()

    # Create mixed workload
    step = [
        MockBatch(op_type="all_reduce"),
        MockBatch(op_type="attention"),
        MockBatch(batch_type="decode"),
    ]

    # Run optimization multiple times with metrics
    for i in range(20):
        strategy.observe_execution(
            batch_size=8,
            latency_ms=100.0,
            metadata={
                "comm_overhead_ratio": 0.3,
                "bubble_time_ratio": 0.15,
            },
        )
        optimized = strategy.optimize(step, stats)

        # Check all optimizations applied
        assert len(optimized) == 3
        if i > 0:  # After first step
            assert strategy._step_count == i + 1

    metrics = strategy.get_metrics()
    assert metrics["step_count"] == 20
    assert metrics["tp_degree"] >= 1
    assert metrics["pp_degree"] >= 1


# Test 12: Communication intensity estimation
def test_comm_intensity_estimation() -> None:
    """Test communication intensity estimation."""
    # Test explicit metadata
    item = MockBatch(metadata={"comm_intensity": 0.7})
    assert ComputeCommOverlapStrategy._estimate_comm_intensity(item) == 0.7

    # Test decode batch
    item = MockBatch(batch_type="decode")
    assert ComputeCommOverlapStrategy._estimate_comm_intensity(item) == 0.2

    # Test prefill batch
    item = MockBatch(batch_type="prefill")
    assert ComputeCommOverlapStrategy._estimate_comm_intensity(item) == 0.5

    # Test communication operations
    item = MockBatch(op_type="all_reduce")
    assert ComputeCommOverlapStrategy._estimate_comm_intensity(item) == 1.0

    # Test compute operations
    item = MockBatch(op_type="attention")
    assert ComputeCommOverlapStrategy._estimate_comm_intensity(item) == 0.1

    # Test default
    item = MockBatch()
    assert ComputeCommOverlapStrategy._estimate_comm_intensity(item) == 0.4


# Test 13: Strategy disabled
def test_strategy_disabled() -> None:
    """Test strategy does nothing when disabled."""
    strategy = ComputeCommOverlapStrategy(enabled=False)
    backend = MockBackend()
    strategy.set_backend(backend)
    stats = RuntimeStats()

    step = [MockBatch(op_type="all_reduce"), MockBatch(op_type="attention")]
    optimized = strategy.optimize(step, stats)

    # When disabled via RuntimeOptimizer, optimize should still run but could be skipped
    # For direct call, it still runs - test the RuntimeOptimizer integration
    assert len(optimized) >= len(step)


# Test 14: RuntimeOptimizer integration
def test_runtime_optimizer_integration() -> None:
    """Test integration with RuntimeOptimizer."""
    backend = MockBackend()
    backend.topology = MockTopology(interconnect="nvlink")

    strategy = ComputeCommOverlapStrategy(
        tp_degree=2,
        pp_degree=2,
    )

    optimizer = RuntimeOptimizer(
        backend=backend,
        strategies=[strategy],
    )

    assert "compute-comm-overlap" in optimizer.list_strategies()

    step = [MockBatch(op_type="attention"), MockBatch(op_type="all_reduce")]
    optimized = optimizer.optimize_step(step)

    # Check optimization was applied
    assert len(optimized) == 2


# Test 15: Throughput improvement calculation
def test_throughput_improvement_calculation() -> None:
    """Test throughput improvement calculation."""
    strategy = ComputeCommOverlapStrategy()
    backend = MockBackend()
    strategy.set_backend(backend)

    # Simulate baseline throughput
    for i in range(10):
        strategy.observe_execution(
            batch_size=8,
            latency_ms=100.0,  # 80 req/s
            metadata={},
        )

    # Simulate improved throughput
    for i in range(10):
        strategy.observe_execution(
            batch_size=8,
            latency_ms=80.0,  # 100 req/s
            metadata={},
        )

    metrics = strategy.get_metrics()

    # Should show improvement
    assert metrics["throughput_improvement"] > 0.0


# Test 16: Communication overhead reduction calculation
def test_comm_overhead_reduction_calculation() -> None:
    """Test communication overhead reduction calculation."""
    strategy = ComputeCommOverlapStrategy()
    backend = MockBackend()
    strategy.set_backend(backend)

    # Simulate baseline overhead
    for i in range(10):
        strategy.observe_execution(
            batch_size=8,
            latency_ms=100.0,
            metadata={"comm_overhead_ratio": 0.4},
        )

    # Simulate reduced overhead
    for i in range(10):
        strategy.observe_execution(
            batch_size=8,
            latency_ms=100.0,
            metadata={"comm_overhead_ratio": 0.25},
        )

    metrics = strategy.get_metrics()

    # Should show reduction
    assert metrics["comm_overhead_reduction"] > 0.0
