"""Tests for analysis modules."""

import pytest

from wafer_core.lib.distributed_traces.analysis import (
    analyze_bandwidth,
    analyze_overlap,
    detect_stragglers,
    analyze_hang_risk,
)
from wafer_core.lib.distributed_traces.analysis.bandwidth import InterconnectType
from wafer_core.lib.distributed_traces.correlation import match_collectives_across_ranks
from wafer_core.lib.distributed_traces.models import (
    Collective,
    CollectiveType,
    ComputeEvent,
    HangRiskLevel,
    RankTimeline,
    TraceSession,
)


class TestBandwidthAnalysis:
    """Tests for bandwidth calculation."""

    def test_calculates_achieved_bandwidth(self, sample_session: TraceSession) -> None:
        """Should calculate achieved bandwidth."""
        result = analyze_bandwidth(sample_session)
        
        assert result.aggregate.achieved_bandwidth_gbps > 0

    def test_calculates_efficiency(self, sample_session: TraceSession) -> None:
        """Should calculate bandwidth efficiency."""
        result = analyze_bandwidth(
            sample_session,
            interconnect=InterconnectType.NVLINK_4,
        )
        
        # Efficiency should be between 0 and 100%
        assert 0 <= result.aggregate.efficiency_percent <= 100

    def test_flags_low_efficiency(self) -> None:
        """Should flag low efficiency collectives."""
        session = TraceSession(name="test")
        t = RankTimeline(rank=0)
        
        # Very slow collective = low efficiency
        t.add_collective(Collective(
            collective_type=CollectiveType.ALL_REDUCE,
            rank=0,
            start_time_ns=0,
            end_time_ns=1_000_000_000,  # 1 second
            message_size_bytes=1000,  # Only 1KB
        ))
        session.add_timeline(t)
        
        result = analyze_bandwidth(
            session,
            efficiency_threshold=50.0,
        )
        
        # Should flag this as low efficiency
        assert len(result.low_efficiency_collectives) == 1


class TestOverlapAnalysis:
    """Tests for communication/compute overlap analysis."""

    def test_detects_overlap(self) -> None:
        """Should detect overlapping compute and communication."""
        session = TraceSession(name="test")
        t = RankTimeline(rank=0)
        
        # Communication from 0 to 10ms
        t.add_collective(Collective(
            collective_type=CollectiveType.ALL_REDUCE,
            rank=0,
            start_time_ns=0,
            end_time_ns=10_000_000,
        ))
        
        # Compute from 5 to 15ms (overlaps with communication)
        t.add_compute_event(ComputeEvent(
            name="matmul",
            start_time_ns=5_000_000,
            end_time_ns=15_000_000,
        ))
        
        session.add_timeline(t)
        
        result = analyze_overlap(session)
        
        assert result.overall.overlapped_time_ns > 0
        assert result.overall.overlap_ratio > 0

    def test_no_overlap_sequential(self) -> None:
        """Should detect zero overlap for sequential operations."""
        session = TraceSession(name="test")
        t = RankTimeline(rank=0)
        
        # Communication from 0 to 10ms
        t.add_collective(Collective(
            collective_type=CollectiveType.ALL_REDUCE,
            rank=0,
            start_time_ns=0,
            end_time_ns=10_000_000,
        ))
        
        # Compute from 10 to 20ms (no overlap)
        t.add_compute_event(ComputeEvent(
            name="matmul",
            start_time_ns=10_000_000,
            end_time_ns=20_000_000,
        ))
        
        session.add_timeline(t)
        
        result = analyze_overlap(session)
        
        assert result.overall.overlapped_time_ns == 0
        assert result.overall.overlap_ratio == 0

    def test_generates_recommendations(self, sample_session: TraceSession) -> None:
        """Should generate overlap improvement recommendations."""
        result = analyze_overlap(sample_session)
        
        # Should have some recommendations
        assert len(result.recommendations) >= 0


class TestStragglerDetection:
    """Tests for straggler detection."""

    def test_detects_straggler(self, sample_session: TraceSession) -> None:
        """Should detect straggler rank."""
        matches = match_collectives_across_ranks(sample_session)
        report = detect_stragglers(sample_session, matches)
        
        # Rank 3 is slowest in our fixture
        assert 3 in report.flagged_ranks or report.rank_scores[3] > report.rank_scores[0]

    def test_computes_total_impact(self, sample_session: TraceSession) -> None:
        """Should compute total straggler impact."""
        matches = match_collectives_across_ranks(sample_session)
        report = detect_stragglers(sample_session, matches)
        
        assert report.total_straggler_impact_ns >= 0

    def test_no_stragglers_equal_timing(self) -> None:
        """Should not flag stragglers when timing is equal."""
        session = TraceSession(name="test")
        
        for rank in range(4):
            t = RankTimeline(rank=rank)
            t.add_collective(Collective(
                collective_type=CollectiveType.ALL_REDUCE,
                rank=rank,
                start_time_ns=1000000,
                end_time_ns=2000000,  # Same duration for all
                sequence_number=1,
            ))
            session.add_timeline(t)
        
        matches = match_collectives_across_ranks(session)
        report = detect_stragglers(session, matches)
        
        # No flagged stragglers when all equal
        assert len(report.flagged_ranks) == 0


class TestHangRiskAnalysis:
    """Tests for hang risk assessment."""

    def test_normal_risk_for_short_collectives(self, sample_session: TraceSession) -> None:
        """Short collectives should have normal risk."""
        report = analyze_hang_risk(sample_session, timeout_sec=600)
        
        # Our fixtures have ms-scale collectives, should be normal
        assert report.risk_distribution[HangRiskLevel.NORMAL] > 0

    def test_critical_risk_for_long_collectives(self) -> None:
        """Collectives approaching timeout should be critical."""
        session = TraceSession(name="test")
        t = RankTimeline(rank=0)
        
        # 350 second collective (>50% of 600s timeout)
        t.add_collective(Collective(
            collective_type=CollectiveType.ALL_REDUCE,
            rank=0,
            start_time_ns=0,
            end_time_ns=350_000_000_000,
        ))
        session.add_timeline(t)
        
        report = analyze_hang_risk(session, timeout_sec=600)
        
        assert report.risk_distribution[HangRiskLevel.CRITICAL] == 1

    def test_generates_recommendations(self) -> None:
        """Should generate recommendations for risky collectives."""
        session = TraceSession(name="test")
        t = RankTimeline(rank=0)
        
        # Warning-level collective (10-50% of timeout)
        t.add_collective(Collective(
            collective_type=CollectiveType.ALL_REDUCE,
            rank=0,
            start_time_ns=0,
            end_time_ns=100_000_000_000,  # 100s = 16.7% of 600s
        ))
        session.add_timeline(t)
        
        report = analyze_hang_risk(session, timeout_sec=600)
        
        assert len(report.recommendations) > 0
