"""Shared fixtures for distributed traces tests."""

import json
from pathlib import Path

import pytest

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


# ── Sample Trace Data ────────────────────────────────────────────────────────


@pytest.fixture
def sample_pytorch_trace() -> dict:
    """Create a sample PyTorch Profiler Chrome trace format."""
    return {
        "traceEvents": [
            # Metadata
            {"ph": "M", "cat": "__metadata", "name": "process_name", "pid": 0, "args": {"name": "Rank 0"}},
            
            # NCCL AllReduce operation
            {
                "ph": "X",
                "cat": "cpu_op",
                "name": "nccl:all_reduce",
                "pid": 0,
                "tid": 0,
                "ts": 1000000,  # 1s in microseconds
                "dur": 5000,    # 5ms
                "args": {
                    "message_size": 4194304,  # 4MB
                    "count": 1048576,
                    "dtype": "float32",
                    "pg_id": "0",
                }
            },
            # NCCL AllGather operation  
            {
                "ph": "X",
                "cat": "cpu_op",
                "name": "nccl:all_gather",
                "pid": 0,
                "tid": 0,
                "ts": 1010000,
                "dur": 3000,
                "args": {
                    "message_size": 2097152,
                    "dtype": "float32",
                }
            },
            # Compute kernel
            {
                "ph": "X",
                "cat": "kernel",
                "name": "volta_fp16_s884gemm_fp16_128x128_nn",
                "pid": 0,
                "tid": 1,
                "ts": 1000000,
                "dur": 4000,
                "args": {
                    "device": 0,
                    "stream": 7,
                    "grid": [128, 1, 1],
                    "block": [256, 1, 1],
                }
            },
        ],
        "metadata": {"test": True}
    }


@pytest.fixture
def sample_nccl_inspector_lines() -> list[str]:
    """Create sample NCCL Inspector JSONL content."""
    events = [
        # Communicator init
        {
            "type": "comm_init",
            "comm_id": "comm_0",
            "nranks": 4,
            "rank": 0,
        },
        # Collective start
        {
            "type": "coll_start",
            "comm_id": "comm_0",
            "seq": 1,
            "op": "allreduce",
            "count": 1048576,
            "datatype": "float32",
            "startNs": 1000000000,
            "algo": "ring",
            "proto": "LL128",
            "nChannels": 8,
        },
        # Collective end
        {
            "type": "coll_end",
            "comm_id": "comm_0",
            "seq": 1,
            "endNs": 1005000000,
        },
        # Another collective
        {
            "type": "coll_start",
            "comm_id": "comm_0",
            "seq": 2,
            "op": "allgather",
            "count": 524288,
            "datatype": "float32",
            "startNs": 1010000000,
        },
        {
            "type": "coll_end",
            "comm_id": "comm_0",
            "seq": 2,
            "endNs": 1013000000,
        },
    ]
    return [json.dumps(e) for e in events]


@pytest.fixture
def sample_rocprofv3_csv_rows() -> list[dict]:
    """Create sample rocprofv3 CSV rows."""
    return [
        {
            "Domain": "RCCL",
            "Function": "rcclAllReduce",
            "Process_Id": "12345",
            "Thread_Id": "1",
            "Correlation_Id": "100",
            "Start_Timestamp": "1000000000",
            "End_Timestamp": "1005000000",
        },
        {
            "Domain": "RCCL",
            "Function": "rcclAllGather",
            "Process_Id": "12345",
            "Thread_Id": "1",
            "Correlation_Id": "101",
            "Start_Timestamp": "1010000000",
            "End_Timestamp": "1013000000",
        },
        {
            "Domain": "HIP",
            "Function": "hipLaunchKernel",
            "Process_Id": "12345",
            "Thread_Id": "1",
            "Correlation_Id": "102",
            "Start_Timestamp": "1000000000",
            "End_Timestamp": "1004000000",
        },
    ]


# ── Model Fixtures ───────────────────────────────────────────────────────────


@pytest.fixture
def sample_collective() -> Collective:
    """Create a sample Collective."""
    return Collective(
        collective_type=CollectiveType.ALL_REDUCE,
        rank=0,
        start_time_ns=1000000000,
        end_time_ns=1005000000,
        message_size_bytes=4194304,
        count=1048576,
        datatype=DataType.FLOAT32,
        process_group=ProcessGroup(comm_id="comm_0", world_size=4),
        sequence_number=1,
        algorithm="ring",
        protocol="LL128",
    )


@pytest.fixture
def sample_timeline(sample_collective: Collective) -> RankTimeline:
    """Create a sample RankTimeline."""
    timeline = RankTimeline(rank=0, device_id=0)
    timeline.add_collective(sample_collective)
    timeline.add_collective(Collective(
        collective_type=CollectiveType.ALL_GATHER,
        rank=0,
        start_time_ns=1010000000,
        end_time_ns=1013000000,
        message_size_bytes=2097152,
        sequence_number=2,
    ))
    return timeline


@pytest.fixture
def sample_session() -> TraceSession:
    """Create a sample TraceSession with multiple ranks."""
    session = TraceSession(name="test_session", world_size=4)
    
    # Create timelines for 4 ranks
    for rank in range(4):
        timeline = RankTimeline(rank=rank, device_id=rank)
        
        # Add collectives with slight timing variation to simulate stragglers
        base_time = 1000000000
        duration_variation = rank * 500000  # Rank 3 is slowest
        
        timeline.add_collective(Collective(
            collective_type=CollectiveType.ALL_REDUCE,
            rank=rank,
            start_time_ns=base_time,
            end_time_ns=base_time + 5000000 + duration_variation,
            message_size_bytes=4194304,
            sequence_number=1,
            process_group=ProcessGroup(comm_id="0", world_size=4),
        ))
        
        timeline.add_collective(Collective(
            collective_type=CollectiveType.ALL_GATHER,
            rank=rank,
            start_time_ns=base_time + 10000000,
            end_time_ns=base_time + 13000000 + duration_variation,
            message_size_bytes=2097152,
            sequence_number=2,
            process_group=ProcessGroup(comm_id="0", world_size=4),
        ))
        
        session.add_timeline(timeline)
    
    return session


# ── File Fixtures ────────────────────────────────────────────────────────────


@pytest.fixture
def pytorch_trace_file(tmp_path: Path, sample_pytorch_trace: dict) -> Path:
    """Create a temporary PyTorch trace file."""
    trace_file = tmp_path / "trace_rank0.json"
    trace_file.write_text(json.dumps(sample_pytorch_trace))
    return trace_file


@pytest.fixture
def nccl_inspector_file(tmp_path: Path, sample_nccl_inspector_lines: list[str]) -> Path:
    """Create a temporary NCCL Inspector file."""
    log_file = tmp_path / "nccl_inspector_rank0.jsonl"
    log_file.write_text("\n".join(sample_nccl_inspector_lines))
    return log_file


@pytest.fixture
def multi_rank_trace_files(tmp_path: Path) -> list[Path]:
    """Create multiple rank trace files."""
    files = []
    
    for rank in range(4):
        trace = {
            "traceEvents": [
                {
                    "ph": "X",
                    "cat": "cpu_op",
                    "name": "nccl:all_reduce",
                    "pid": rank,
                    "tid": 0,
                    "ts": 1000000 + rank * 100,  # Slight timing offset
                    "dur": 5000 + rank * 500,     # Rank 3 slowest
                    "args": {
                        "message_size": 4194304,
                        "sequence": 1,
                    }
                },
                {
                    "ph": "X",
                    "cat": "cpu_op",
                    "name": "nccl:all_gather",
                    "pid": rank,
                    "tid": 0,
                    "ts": 1010000,
                    "dur": 3000 + rank * 300,
                    "args": {
                        "message_size": 2097152,
                        "sequence": 2,
                    }
                },
            ]
        }
        
        trace_file = tmp_path / f"trace_rank{rank}.json"
        trace_file.write_text(json.dumps(trace))
        files.append(trace_file)
    
    return files
