"""NCCL Inspector JSONL parsing.
Parses trace logs generated by NCCL Inspector plugin when running with
NCCL_PROFILER_PLUGIN=inspector environment variable.
NCCL Inspector provides per-collective timing with microsecond precision,
algorithm selection, and channel utilization information.
Format Reference:
- https://developer.nvidia.com/blog/enhancing-communication-observability-of-ai-workloads-with-nccl-inspector/
"""
import gzip
import json
from pathlib import Path
from typing import Optional
from wafer_core.lib.distributed_traces.models.collective import (
    Collective,
    CollectiveType,
    DataType,
    ProcessGroup,
)
from wafer_core.lib.distributed_traces.models.rank_timeline import RankTimeline
def parse_nccl_inspector_file(
    file_path: str | Path,
    rank: Optional[int] = None,
) -> tuple[RankTimeline | None, str | None]:
    """Parse an NCCL Inspector JSONL file.
    Args:
        file_path: Path to the JSONL or .log.gz file
        rank: Rank number (if known). If None, will try to infer.
    Returns:
        Tuple of (RankTimeline, error). One will be None.
    """
    path = Path(file_path)
    if not path.exists():
        return None, f"File not found: {file_path}"
    try:
        lines = []
        if path.suffix == ".gz":
            with gzip.open(path, "rt", encoding="utf-8") as f:
                lines = f.readlines()
        else:
            with open(path, encoding="utf-8") as f:
                lines = f.readlines()
        return parse_nccl_inspector(lines, rank, str(path))
    except Exception as e:
        return None, f"Failed to parse NCCL Inspector log: {e}"
def parse_nccl_inspector(
    lines: list[str],
    rank: Optional[int] = None,
    filename: Optional[str] = None,
) -> tuple[RankTimeline | None, str | None]:
    """Parse NCCL Inspector log lines.
    NCCL Inspector outputs JSONL where each line is a JSON object
    representing one event. Event types include:
    - Communicator creation/destruction
    - Collective start/end
    - Kernel channel information
    Args:
        lines: Lines from the JSONL file
        rank: Rank number. If None, will infer from events.
        filename: Original filename for rank inference
    Returns:
        Tuple of (RankTimeline, error). One will be None.
    """

    ongoing: dict[tuple[str, int], dict] = {}  # (comm_id, seq) -> start event
    collectives: list[Collective] = []
    communicators: dict[str, ProcessGroup] = {}
    inferred_rank = rank
    for line in lines:
        line = line.strip()
        if not line:
            continue
        try:
            event = json.loads(line)
        except json.JSONDecodeError:
            continue
        event_type = event.get("type", event.get("event", ""))
        if event_type in ["comm_init", "communicator", "ncclCommInitRank"]:
            comm_id = str(event.get("comm_id", event.get("commId", "")))
            nranks = event.get("nranks", event.get("nRanks", 0))
            comm_rank = event.get("rank", event.get("commRank"))
            if comm_rank is not None and inferred_rank is None:
                inferred_rank = comm_rank
            communicators[comm_id] = ProcessGroup(
                comm_id=comm_id,
                world_size=nranks,
            )
            continue
        if event_type in ["coll_start", "collective_start", "start"]:
            comm_id = str(event.get("comm_id", event.get("commId", "")))
            seq = event.get("seq", event.get("seqNum", event.get("sequence", 0)))
            key = (comm_id, seq)
            ongoing[key] = event
            if inferred_rank is None:
                inferred_rank = event.get("rank", event.get("commRank"))
            continue
        if event_type in ["coll_end", "collective_end", "end"]:
            comm_id = str(event.get("comm_id", event.get("commId", "")))
            seq = event.get("seq", event.get("seqNum", event.get("sequence", 0)))
            key = (comm_id, seq)
            start_event = ongoing.pop(key, None)
            if start_event is None:
                # Orphan end event - try to create collective with available info
                start_event = event
            collective = _create_collective_from_events(
                start_event,
                event,
                communicators.get(comm_id),
                inferred_rank or 0,
            )
            if collective is not None:
                collectives.append(collective)
            continue
        if event_type in ["collective", "coll"] or "startNs" in event and "endNs" in event:
            comm_id = str(event.get("comm_id", event.get("commId", "")))
            collective = _create_collective_from_events(
                event,
                event,
                communicators.get(comm_id),
                inferred_rank or 0,
            )
            if collective is not None:
                collectives.append(collective)
    if inferred_rank is None:
        inferred_rank = _infer_rank_from_filename(filename) if filename else 0
    timeline = RankTimeline(rank=inferred_rank)
    for c in collectives:
        timeline.add_collective(c)
    return timeline, None
def _infer_rank_from_filename(filename: str) -> int:
    """Try to infer rank from filename."""
    import re
    match = re.search(r"rank[_\-]?(\d+)", filename.lower())
    if match:
        return int(match.group(1))
    match = re.search(r"_(\d+)\.(jsonl|log)", filename.lower())
    if match:
        return int(match.group(1))
    return 0
def _create_collective_from_events(
    start_event: dict,
    end_event: dict,
    process_group: Optional[ProcessGroup],
    rank: int,
) -> Optional[Collective]:
    """Create a Collective from start and end events."""
    op = start_event.get("op", start_event.get("operation", start_event.get("coll_type", "")))
    collective_type = CollectiveType.from_string(op)
    # Field names like "startNs" indicate nanoseconds, trust the format
    start_ns = start_event.get("startNs", start_event.get("start_ns", start_event.get("timestamp", 0)))
    end_ns = end_event.get("endNs", end_event.get("end_ns", end_event.get("timestamp", 0)))
    # Only convert from microseconds if explicitly using microsecond field names
    if "startUs" in start_event or "start_us" in start_event:
        start_ns = int(start_event.get("startUs", start_event.get("start_us", 0)) * 1000)
        end_ns = int(end_event.get("endUs", end_event.get("end_us", 0)) * 1000)
    if end_ns <= start_ns:
        # Duration might be provided separately
        duration = start_event.get("duration", start_event.get("dur", 0))
        if duration > 0:
            if duration < 1e6:  # Likely microseconds
                duration = int(duration * 1000)
            end_ns = start_ns + duration
        else:
            return None
    count = start_event.get("count", start_event.get("nelem", 0))
    dtype_str = start_event.get("datatype", start_event.get("dtype", "float32"))
    datatype = DataType.from_string(dtype_str)
    message_size = count * datatype.size_bytes()
    # Override with explicit size if available
    if "bytes" in start_event:
        message_size = start_event["bytes"]
    elif "message_size" in start_event:
        message_size = start_event["message_size"]
    algorithm = start_event.get("algo", start_event.get("algorithm"))
    protocol = start_event.get("proto", start_event.get("protocol"))
    num_channels = start_event.get("nChannels", start_event.get("channels"))
    seq = start_event.get("seq", start_event.get("seqNum", start_event.get("sequence")))
    return Collective(
        collective_type=collective_type,
        rank=rank,
        start_time_ns=int(start_ns),
        end_time_ns=int(end_ns),
        message_size_bytes=message_size,
        count=count,
        datatype=datatype,
        process_group=process_group,
        sequence_number=seq,
        algorithm=algorithm,
        protocol=protocol,
        num_channels=num_channels,
        extra={**start_event, **end_event},
    )
