"""Tests for the main distributed traces API."""

from pathlib import Path

import pytest

from wafer_core.lib.distributed_traces import (
    AnalyzerConfig,
    DistributedTracesAnalyzer,
    analyze_distributed_traces,
)
from wafer_core.lib.distributed_traces.analysis.bandwidth import InterconnectType


class TestDistributedTracesAnalyzer:
    """Tests for the main analyzer class."""

    def test_load_single_trace(self, pytorch_trace_file: Path) -> None:
        """Should load a single trace file."""
        analyzer = DistributedTracesAnalyzer()
        session, error = analyzer.load_traces([str(pytorch_trace_file)])
        
        assert error is None
        assert session is not None
        assert session.num_ranks == 1

    def test_load_multiple_traces(self, multi_rank_trace_files: list[Path]) -> None:
        """Should load multiple rank traces."""
        analyzer = DistributedTracesAnalyzer()
        session, error = analyzer.load_traces([str(f) for f in multi_rank_trace_files])
        
        assert error is None
        assert session is not None
        assert session.num_ranks == 4

    def test_analyze_session(self, multi_rank_trace_files: list[Path]) -> None:
        """Should perform full analysis."""
        analyzer = DistributedTracesAnalyzer()
        session, _ = analyzer.load_traces([str(f) for f in multi_rank_trace_files])
        
        assert session is not None
        result = analyzer.analyze(session)
        
        # Check all components present
        assert result.aggregate_bandwidth is not None
        assert result.overlap_metrics is not None
        assert result.straggler_report is not None
        assert result.hang_risk_report is not None
        assert result.summary is not None

    def test_analyze_with_config(self, multi_rank_trace_files: list[Path]) -> None:
        """Should use custom configuration."""
        config = AnalyzerConfig(
            interconnect=InterconnectType.NVLINK_4,
            timeout_sec=300,
        )
        analyzer = DistributedTracesAnalyzer(config)
        
        session, _ = analyzer.load_traces([str(f) for f in multi_rank_trace_files])
        assert session is not None
        
        result = analyzer.analyze(session)
        
        # Theoretical bandwidth should match NVLink 4
        assert result.aggregate_bandwidth.theoretical_bandwidth_gbps == 900.0

    def test_export_json(self, multi_rank_trace_files: list[Path], tmp_path: Path) -> None:
        """Should export to JSON."""
        analyzer = DistributedTracesAnalyzer()
        session, _ = analyzer.load_traces([str(f) for f in multi_rank_trace_files])
        assert session is not None
        
        result = analyzer.analyze(session)
        
        output_file = tmp_path / "analysis.json"
        json_str = analyzer.export_json(result, session, str(output_file))
        
        assert output_file.exists()
        assert "aggregate_bandwidth" in json_str
        assert "overlap_metrics" in json_str

    def test_export_csv(self, multi_rank_trace_files: list[Path], tmp_path: Path) -> None:
        """Should export to CSV."""
        analyzer = DistributedTracesAnalyzer()
        session, _ = analyzer.load_traces([str(f) for f in multi_rank_trace_files])
        assert session is not None
        
        output_file = tmp_path / "collectives.csv"
        csv_str = analyzer.export_csv(session, str(output_file))
        
        assert output_file.exists()
        assert "rank,collective_type" in csv_str

    def test_export_timeline(self, multi_rank_trace_files: list[Path]) -> None:
        """Should export timeline data."""
        analyzer = DistributedTracesAnalyzer()
        session, _ = analyzer.load_traces([str(f) for f in multi_rank_trace_files])
        assert session is not None
        
        timeline_json = analyzer.export_timeline(session)
        
        assert "metadata" in timeline_json
        assert "ranks" in timeline_json

    def test_compare_sessions(self, multi_rank_trace_files: list[Path]) -> None:
        """Should compare two sessions."""
        analyzer = DistributedTracesAnalyzer()
        
        # Load same traces as baseline and comparison (for testing structure)
        baseline, _ = analyzer.load_traces([str(f) for f in multi_rank_trace_files])
        comparison, _ = analyzer.load_traces([str(f) for f in multi_rank_trace_files])
        
        assert baseline is not None
        assert comparison is not None
        
        result = analyzer.compare_sessions(baseline, comparison)
        
        # Same sessions should have zero delta
        assert result.bandwidth_delta == 0
        assert result.overlap_delta == 0

    def test_load_error_handling(self, tmp_path: Path) -> None:
        """Should handle errors gracefully."""
        analyzer = DistributedTracesAnalyzer()
        
        # Non-existent file
        session, error = analyzer.load_traces([str(tmp_path / "nonexistent.json")])
        
        assert error is not None
        assert session is None

    def test_load_empty_list(self) -> None:
        """Should handle empty file list."""
        analyzer = DistributedTracesAnalyzer()
        session, error = analyzer.load_traces([])
        
        assert error is not None
        assert "No trace files" in error


class TestConvenienceFunction:
    """Tests for the analyze_distributed_traces convenience function."""

    def test_analyze_distributed_traces(self, multi_rank_trace_files: list[Path]) -> None:
        """Should work with convenience function."""
        result, error = analyze_distributed_traces(
            [str(f) for f in multi_rank_trace_files]
        )
        
        assert error is None
        assert result is not None
        assert result.summary is not None

    def test_with_config(self, multi_rank_trace_files: list[Path]) -> None:
        """Should accept config parameter."""
        config = AnalyzerConfig(timeout_sec=300)
        
        result, error = analyze_distributed_traces(
            [str(f) for f in multi_rank_trace_files],
            config=config,
        )
        
        assert error is None
        assert result is not None
