"""Tests for distributed traces parsers."""

import json
from pathlib import Path

import pytest

from wafer_core.lib.distributed_traces.models import CollectiveType
from wafer_core.lib.distributed_traces.parsers import (
    parse_pytorch_trace,
    parse_pytorch_trace_file,
    parse_nccl_inspector,
    parse_nccl_inspector_file,
    parse_rocprofv3,
)


class TestPyTorchTraceParser:
    """Tests for PyTorch Profiler trace parsing."""

    def test_parse_file_success(self, pytorch_trace_file: Path) -> None:
        """Should parse PyTorch trace file successfully."""
        timeline, error = parse_pytorch_trace_file(pytorch_trace_file, rank=0)
        
        assert error is None
        assert timeline is not None
        assert timeline.rank == 0

    def test_parse_extracts_collectives(self, sample_pytorch_trace: dict) -> None:
        """Should extract NCCL collectives from trace."""
        timeline, error = parse_pytorch_trace(sample_pytorch_trace, rank=0)
        
        assert error is None
        assert timeline is not None
        assert len(timeline.collectives) == 2

    def test_parse_collective_types(self, sample_pytorch_trace: dict) -> None:
        """Should correctly identify collective types."""
        timeline, _ = parse_pytorch_trace(sample_pytorch_trace, rank=0)
        
        assert timeline is not None
        types = [c.collective_type for c in timeline.collectives]
        assert CollectiveType.ALL_REDUCE in types
        assert CollectiveType.ALL_GATHER in types

    def test_parse_extracts_compute_events(self, sample_pytorch_trace: dict) -> None:
        """Should extract compute kernels."""
        timeline, _ = parse_pytorch_trace(sample_pytorch_trace, rank=0)
        
        assert timeline is not None
        assert len(timeline.compute_events) >= 1
        assert "gemm" in timeline.compute_events[0].name.lower()

    def test_parse_message_size(self, sample_pytorch_trace: dict) -> None:
        """Should extract message size."""
        timeline, _ = parse_pytorch_trace(sample_pytorch_trace, rank=0)
        
        assert timeline is not None
        all_reduce = next(
            c for c in timeline.collectives 
            if c.collective_type == CollectiveType.ALL_REDUCE
        )
        assert all_reduce.message_size_bytes == 4194304

    def test_parse_file_not_found(self, tmp_path: Path) -> None:
        """Should return error for missing file."""
        _, error = parse_pytorch_trace_file(tmp_path / "nonexistent.json")
        
        assert error is not None
        assert "not found" in error.lower()

    def test_parse_invalid_json(self, tmp_path: Path) -> None:
        """Should return error for invalid JSON."""
        bad_file = tmp_path / "bad.json"
        bad_file.write_text("not valid json")
        
        _, error = parse_pytorch_trace_file(bad_file)
        
        assert error is not None

    def test_infer_rank_from_filename(self, tmp_path: Path, sample_pytorch_trace: dict) -> None:
        """Should infer rank from filename pattern."""
        trace_file = tmp_path / "trace_rank3.json"
        trace_file.write_text(json.dumps(sample_pytorch_trace))
        
        timeline, _ = parse_pytorch_trace_file(trace_file)
        
        assert timeline is not None
        assert timeline.rank == 3


class TestNCCLInspectorParser:
    """Tests for NCCL Inspector parsing."""

    def test_parse_file_success(self, nccl_inspector_file: Path) -> None:
        """Should parse NCCL Inspector file successfully."""
        timeline, error = parse_nccl_inspector_file(nccl_inspector_file, rank=0)
        
        assert error is None
        assert timeline is not None

    def test_parse_extracts_collectives(self, sample_nccl_inspector_lines: list[str]) -> None:
        """Should extract collectives from JSONL."""
        timeline, error = parse_nccl_inspector(sample_nccl_inspector_lines, rank=0)
        
        assert error is None
        assert timeline is not None
        assert len(timeline.collectives) == 2

    def test_parse_collective_metadata(self, sample_nccl_inspector_lines: list[str]) -> None:
        """Should extract collective metadata."""
        timeline, _ = parse_nccl_inspector(sample_nccl_inspector_lines, rank=0)
        
        assert timeline is not None
        all_reduce = next(
            c for c in timeline.collectives 
            if c.collective_type == CollectiveType.ALL_REDUCE
        )
        assert all_reduce.algorithm == "ring"
        assert all_reduce.protocol == "LL128"
        assert all_reduce.num_channels == 8

    def test_parse_timing(self, sample_nccl_inspector_lines: list[str]) -> None:
        """Should compute correct timing from start/end events."""
        timeline, _ = parse_nccl_inspector(sample_nccl_inspector_lines, rank=0)
        
        assert timeline is not None
        all_reduce = next(
            c for c in timeline.collectives 
            if c.collective_type == CollectiveType.ALL_REDUCE
        )
        # endNs (1005000000) - startNs (1000000000) = 5000000 ns = 5ms
        assert all_reduce.duration_ns == 5000000


class TestRocprofv3Parser:
    """Tests for rocprofv3 CSV/JSON parsing."""

    def test_parse_csv_rows(self, sample_rocprofv3_csv_rows: list[dict]) -> None:
        """Should parse CSV rows into timeline."""
        timeline, error = parse_rocprofv3(sample_rocprofv3_csv_rows, rank=0, is_csv=True)
        
        assert error is None
        assert timeline is not None

    def test_parse_identifies_rccl(self, sample_rocprofv3_csv_rows: list[dict]) -> None:
        """Should identify RCCL operations."""
        timeline, _ = parse_rocprofv3(sample_rocprofv3_csv_rows, rank=0, is_csv=True)
        
        assert timeline is not None
        assert len(timeline.collectives) == 2

    def test_parse_collective_types(self, sample_rocprofv3_csv_rows: list[dict]) -> None:
        """Should identify collective types from function names."""
        timeline, _ = parse_rocprofv3(sample_rocprofv3_csv_rows, rank=0, is_csv=True)
        
        assert timeline is not None
        types = {c.collective_type for c in timeline.collectives}
        assert CollectiveType.ALL_REDUCE in types
        assert CollectiveType.ALL_GATHER in types

    def test_parse_filters_hip_kernels(self, sample_rocprofv3_csv_rows: list[dict]) -> None:
        """Should filter HIP kernel launches from collectives."""
        timeline, _ = parse_rocprofv3(sample_rocprofv3_csv_rows, rank=0, is_csv=True)
        
        assert timeline is not None
        # HIP kernel should be in compute events, not collectives
        assert len(timeline.collectives) == 2
