"""Tests for collective operations benchmark.

Verifies benchmark functionality and output format.
"""

from __future__ import annotations

import tempfile
from pathlib import Path

import pytest

from sagellm_comm import GlooBackend
from sagellm_comm.benchmark import BenchmarkResult, CollectiveBenchmark


class TestCollectiveBenchmark:
    """Test CollectiveBenchmark class."""

    @pytest.fixture
    def backend(self) -> GlooBackend:
        """Initialize Gloo backend for testing."""
        backend = GlooBackend()
        backend.init(rank=0, world_size=1, master_addr="localhost", master_port=29500)
        yield backend
        backend.destroy()

    def test_benchmark_result(self):
        """Test BenchmarkResult dataclass."""
        result = BenchmarkResult(
            bytes=4096,
            latency_ms=1.5,
            bandwidth_gbps=2.5,
            algo_id="ring",
            backend_kind="gloo",
            topology_source="env",
            world_size=2,
            op_type="all_reduce",
            data_type="float32",
            warmup_iters=10,
            benchmark_iters=100,
            selection_rule_id="rule_1",
            selection_reason="test reason",
            timestamp="2024-01-01 00:00:00",
        )

        assert result.bytes == 4096
        assert result.latency_ms == 1.5
        assert result.bandwidth_gbps == 2.5
        assert result.algo_id == "ring"
        assert result.backend_kind == "gloo"

    def test_benchmark_initialization(self, backend: GlooBackend):
        """Test CollectiveBenchmark initialization."""
        benchmark = CollectiveBenchmark(backend, warmup_iters=5, benchmark_iters=10)

        assert benchmark.backend == backend
        assert benchmark.warmup_iters == 5
        assert benchmark.benchmark_iters == 10
        assert benchmark.algo_id == "auto"

    def test_run_all_reduce_basic(self, backend: GlooBackend):
        """Test basic all-reduce benchmark."""
        benchmark = CollectiveBenchmark(backend, warmup_iters=1, benchmark_iters=2)

        # Small message sizes for fast testing
        results = benchmark.run_all_reduce(message_sizes=[10, 100])

        assert len(results) == 2

        # Check first result
        result = results[0]
        assert result.bytes == 10 * 4  # 10 float32 elements = 40 bytes
        assert result.latency_ms > 0
        assert result.bandwidth_gbps > 0
        assert result.algo_id in {"ring", "tree", "hybrid"}
        assert result.backend_kind == "gloo"
        assert result.selection_rule_id
        assert result.selection_reason
        assert result.world_size == 1
        assert result.op_type == "all_reduce"
        assert result.data_type == "float32"
        assert result.warmup_iters == 1
        assert result.benchmark_iters == 2

    def test_export_json(self, backend: GlooBackend):
        """Test JSON export."""
        benchmark = CollectiveBenchmark(backend, warmup_iters=1, benchmark_iters=1)
        results = benchmark.run_all_reduce(message_sizes=[10])

        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "results.json"
            benchmark.export_json(results, output_path)

            assert output_path.exists()

            # Verify JSON content
            import json

            with open(output_path) as f:
                data = json.load(f)

            assert len(data) == 1
            assert data[0]["bytes"] == 40
            assert "latency_ms" in data[0]
            assert "bandwidth_gbps" in data[0]

    def test_export_csv(self, backend: GlooBackend):
        """Test CSV export."""
        benchmark = CollectiveBenchmark(backend, warmup_iters=1, benchmark_iters=1)
        results = benchmark.run_all_reduce(message_sizes=[10, 20])

        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "results.csv"
            benchmark.export_csv(results, output_path)

            assert output_path.exists()

            # Verify CSV content
            import csv

            with open(output_path) as f:
                reader = csv.DictReader(f)
                rows = list(reader)

            assert len(rows) == 2
            assert rows[0]["bytes"] == "40"  # CSV reads as strings
            assert "latency_ms" in rows[0]
            assert "bandwidth_gbps" in rows[0]

    def test_export_empty_results_error(self):
        """Test that exporting empty results raises error."""
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "results.csv"

            with pytest.raises(ValueError, match="Cannot export empty results"):
                CollectiveBenchmark.export_csv([], output_path)
