"""Tests for distributed traces data models."""

import pytest

from wafer_core.lib.distributed_traces.models import (
    Collective,
    CollectiveType,
    DataType,
    ProcessGroup,
    RankTimeline,
    TraceSession,
    TraceFormat,
)


class TestCollectiveType:
    """Tests for CollectiveType enum."""

    def test_from_string_nccl_prefix(self) -> None:
        """Should parse nccl: prefixed names."""
        assert CollectiveType.from_string("nccl:all_reduce") == CollectiveType.ALL_REDUCE
        assert CollectiveType.from_string("nccl:allgather") == CollectiveType.ALL_GATHER

    def test_from_string_rccl_prefix(self) -> None:
        """Should parse rccl: prefixed names."""
        assert CollectiveType.from_string("rccl:all_reduce") == CollectiveType.ALL_REDUCE
        assert CollectiveType.from_string("rccl:reducescatter") == CollectiveType.REDUCE_SCATTER

    def test_from_string_camel_case(self) -> None:
        """Should parse camelCase names."""
        assert CollectiveType.from_string("ncclAllReduce") == CollectiveType.ALL_REDUCE
        assert CollectiveType.from_string("rcclAllGather") == CollectiveType.ALL_GATHER

    def test_from_string_unknown(self) -> None:
        """Should return UNKNOWN for unrecognized types."""
        assert CollectiveType.from_string("unknown_op") == CollectiveType.UNKNOWN


class TestDataType:
    """Tests for DataType enum."""

    def test_size_bytes(self) -> None:
        """Should return correct byte sizes."""
        assert DataType.FLOAT32.size_bytes() == 4
        assert DataType.FLOAT16.size_bytes() == 2
        assert DataType.FLOAT64.size_bytes() == 8
        assert DataType.INT8.size_bytes() == 1

    def test_from_string(self) -> None:
        """Should parse data type strings."""
        assert DataType.from_string("float32") == DataType.FLOAT32
        assert DataType.from_string("fp16") == DataType.FLOAT16
        assert DataType.from_string("bfloat16") == DataType.BFLOAT16


class TestCollective:
    """Tests for Collective dataclass."""

    def test_immutable(self, sample_collective: Collective) -> None:
        """Collective should be immutable (frozen)."""
        with pytest.raises(Exception):  # FrozenInstanceError
            sample_collective.rank = 1  # type: ignore

    def test_duration_properties(self, sample_collective: Collective) -> None:
        """Should calculate duration correctly."""
        assert sample_collective.duration_ns == 5000000
        assert sample_collective.duration_ms == 5.0
        assert sample_collective.duration_us == 5000.0

    def test_validation_negative_rank(self) -> None:
        """Should reject negative rank."""
        with pytest.raises(AssertionError):
            Collective(
                collective_type=CollectiveType.ALL_REDUCE,
                rank=-1,
                start_time_ns=0,
                end_time_ns=1000,
            )

    def test_validation_end_before_start(self) -> None:
        """Should reject end time before start time."""
        with pytest.raises(AssertionError):
            Collective(
                collective_type=CollectiveType.ALL_REDUCE,
                rank=0,
                start_time_ns=1000,
                end_time_ns=500,
            )


class TestRankTimeline:
    """Tests for RankTimeline."""

    def test_add_collective(self, sample_collective: Collective) -> None:
        """Should add collectives to timeline."""
        timeline = RankTimeline(rank=0)
        timeline.add_collective(sample_collective)
        
        assert len(timeline.collectives) == 1
        assert timeline.collectives[0] == sample_collective

    def test_time_properties(self, sample_timeline: RankTimeline) -> None:
        """Should calculate time range correctly."""
        assert sample_timeline.start_time_ns == 1000000000
        assert sample_timeline.end_time_ns == 1013000000
        assert sample_timeline.duration_ns == 13000000

    def test_total_collective_time(self, sample_timeline: RankTimeline) -> None:
        """Should sum collective durations."""
        # 5ms + 3ms = 8ms = 8000000ns
        assert sample_timeline.total_collective_time_ns == 8000000

    def test_get_collectives_by_type(self, sample_timeline: RankTimeline) -> None:
        """Should filter collectives by type."""
        all_reduces = sample_timeline.get_collectives_by_type(CollectiveType.ALL_REDUCE)
        assert len(all_reduces) == 1
        assert all_reduces[0].collective_type == CollectiveType.ALL_REDUCE


class TestTraceSession:
    """Tests for TraceSession."""

    def test_add_timeline(self, sample_timeline: RankTimeline) -> None:
        """Should add timelines to session."""
        session = TraceSession(name="test")
        session.add_timeline(sample_timeline)
        
        assert session.num_ranks == 1
        assert 0 in session.ranks
        assert session.get_timeline(0) == sample_timeline

    def test_total_collectives(self, sample_session: TraceSession) -> None:
        """Should count total collectives across ranks."""
        # 4 ranks × 2 collectives = 8
        assert sample_session.total_collectives == 8

    def test_get_all_collectives(self, sample_session: TraceSession) -> None:
        """Should return all collectives from all ranks."""
        collectives = sample_session.get_all_collectives()
        assert len(collectives) == 8

    def test_get_collectives_by_type(self, sample_session: TraceSession) -> None:
        """Should filter by type across all ranks."""
        all_reduces = sample_session.get_collectives_by_type(CollectiveType.ALL_REDUCE)
        assert len(all_reduces) == 4  # One per rank


class TestTraceFormat:
    """Tests for TraceFormat auto-detection."""

    def test_detect_pytorch_json(self, tmp_path) -> None:
        """Should detect PyTorch Profiler JSON."""
        trace_file = tmp_path / "trace.json"
        trace_file.write_text('{"traceEvents": []}')
        
        assert TraceFormat.detect_format(str(trace_file)) == TraceFormat.PYTORCH_PROFILER

    def test_detect_nccl_inspector(self, tmp_path) -> None:
        """Should detect NCCL Inspector JSONL."""
        log_file = tmp_path / "inspector.jsonl"
        log_file.write_text('{"op": "allreduce", "comm_id": "0"}')
        
        assert TraceFormat.detect_format(str(log_file)) == TraceFormat.NCCL_INSPECTOR

    def test_detect_nsys(self, tmp_path) -> None:
        """Should detect Nsight Systems."""
        nsys_file = tmp_path / "report.nsys-rep"
        nsys_file.write_text("")  # Extension is key
        
        assert TraceFormat.detect_format(str(nsys_file)) == TraceFormat.NSYS

    def test_detect_rocprofv3_csv(self, tmp_path) -> None:
        """Should detect rocprofv3 CSV."""
        csv_file = tmp_path / "results.csv"
        csv_file.write_text("Domain,Function,RCCL")
        
        assert TraceFormat.detect_format(str(csv_file)) == TraceFormat.ROCPROFV3
