"""Tests for export modules."""

import json
from pathlib import Path

import pytest

from wafer_core.lib.distributed_traces import DistributedTracesAnalyzer
from wafer_core.lib.distributed_traces.export import (
    export_analysis_json,
    export_collectives_csv,
    export_session_json,
    export_timeline_json,
    export_perfetto_format,
)


class TestJsonExport:
    """Tests for JSON export functionality."""

    def test_export_session_json(self, sample_session) -> None:
        """Should export session to JSON."""
        json_str = export_session_json(sample_session)
        
        data = json.loads(json_str)
        assert data["name"] == "test_session"
        assert data["world_size"] == 4
        assert "timelines" in data

    def test_export_session_without_events(self, sample_session) -> None:
        """Should support excluding events."""
        json_str = export_session_json(sample_session, include_events=False)
        
        data = json.loads(json_str)
        assert "timelines" not in data

    def test_export_to_file(self, sample_session, tmp_path: Path) -> None:
        """Should write to file when path provided."""
        output_file = tmp_path / "session.json"
        export_session_json(sample_session, output_path=str(output_file))
        
        assert output_file.exists()
        data = json.loads(output_file.read_text())
        assert data["name"] == "test_session"


class TestCsvExport:
    """Tests for CSV export functionality."""

    def test_export_collectives_csv(self, sample_session) -> None:
        """Should export collectives to CSV."""
        csv_str = export_collectives_csv(sample_session)
        
        lines = csv_str.strip().split("\n")
        assert len(lines) >= 2  # Header + at least one data row
        
        # Check header
        header = lines[0]
        assert "rank" in header
        assert "collective_type" in header
        assert "duration_ms" in header

    def test_csv_data_rows(self, sample_session) -> None:
        """Should include correct data rows."""
        csv_str = export_collectives_csv(sample_session)
        
        lines = csv_str.strip().split("\n")
        # 4 ranks × 2 collectives + 1 header = 9 lines
        assert len(lines) == 9

    def test_export_to_file(self, sample_session, tmp_path: Path) -> None:
        """Should write to file when path provided."""
        output_file = tmp_path / "collectives.csv"
        export_collectives_csv(sample_session, output_path=str(output_file))
        
        assert output_file.exists()


class TestTimelineExport:
    """Tests for timeline export functionality."""

    def test_export_timeline_json(self, sample_session) -> None:
        """Should export timeline data for visualization."""
        json_str = export_timeline_json(sample_session)
        
        data = json.loads(json_str)
        assert "metadata" in data
        assert "ranks" in data
        assert len(data["ranks"]) == 4

    def test_timeline_has_events(self, sample_session) -> None:
        """Should include events in timeline."""
        json_str = export_timeline_json(sample_session)
        
        data = json.loads(json_str)
        assert len(data["ranks"][0]["events"]) > 0

    def test_timeline_alignment_lines(self, sample_session) -> None:
        """Should include alignment lines when matches provided."""
        from wafer_core.lib.distributed_traces.correlation import match_collectives_across_ranks
        
        matches = match_collectives_across_ranks(sample_session)
        json_str = export_timeline_json(sample_session, matches=matches)
        
        data = json.loads(json_str)
        assert "alignment_lines" in data


class TestPerfettoExport:
    """Tests for Perfetto format export."""

    def test_export_perfetto_format(self, sample_session) -> None:
        """Should export to Perfetto-compatible format."""
        json_str = export_perfetto_format(sample_session)
        
        data = json.loads(json_str)
        assert "traceEvents" in data
        assert "metadata" in data

    def test_perfetto_has_process_names(self, sample_session) -> None:
        """Should include process name metadata."""
        json_str = export_perfetto_format(sample_session)
        
        data = json.loads(json_str)
        process_names = [
            e for e in data["traceEvents"]
            if e.get("name") == "process_name"
        ]
        assert len(process_names) == 4  # One per rank

    def test_perfetto_event_format(self, sample_session) -> None:
        """Should use correct Chrome trace event format."""
        json_str = export_perfetto_format(sample_session)
        
        data = json.loads(json_str)
        events = [e for e in data["traceEvents"] if e.get("ph") == "X"]
        
        assert len(events) > 0
        event = events[0]
        assert "ts" in event  # Timestamp in microseconds
        assert "dur" in event  # Duration
        assert "cat" in event  # Category
        assert "name" in event
