"""PyTorch Profiler Chrome trace JSON parsing.
Parses trace files generated by torch.profiler.profile() with export_chrome_trace().
Extracts NCCL collective events by matching kernel names and operator names.
Format Reference:
- Chrome trace format: https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU
- PyTorch profiler: https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html
"""
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 ComputeEvent, RankTimeline
def parse_pytorch_trace_file(
    file_path: str | Path,
    rank: Optional[int] = None,
) -> tuple[RankTimeline | None, str | None]:
    """Parse a PyTorch profiler trace file.
    Args:
        file_path: Path to the trace JSON or JSON.gz file
        rank: Rank number (if known). If None, will try to infer from filename.
    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:
        if path.suffix == ".gz":
            with gzip.open(path, "rt", encoding="utf-8") as f:
                data = json.load(f)
        else:
            with open(path, encoding="utf-8") as f:
                data = json.load(f)
        return parse_pytorch_trace(data, rank, str(path))
    except json.JSONDecodeError as e:
        return None, f"JSON parse error: {e}"
    except Exception as e:
        return None, f"Failed to parse trace: {e}"
def parse_pytorch_trace(
    data: dict,
    rank: Optional[int] = None,
    filename: Optional[str] = None,
) -> tuple[RankTimeline | None, str | None]:
    """Parse PyTorch profiler trace data.
    Args:
        data: Parsed JSON data from Chrome trace format
        rank: Rank number. If None, will try to infer.
        filename: Original filename for rank inference
    Returns:
        Tuple of (RankTimeline, error). One will be None.
    """
    # Try to infer rank from filename if not provided
    if rank is None:
        rank = _infer_rank_from_filename(filename) if filename else 0
    timeline = RankTimeline(rank=rank)
    events = data.get("traceEvents", data.get("trace_events", []))
    if not events:
        return None, "No trace events found in file"
    correlation_map: dict[int, dict] = {}
    # First pass: collect all events and build correlation map
    for event in events:
        if not isinstance(event, dict):
            continue
        corr_id = event.get("args", {}).get("correlation", event.get("args", {}).get("correlationId"))
        if corr_id is not None:
            if corr_id not in correlation_map:
                correlation_map[corr_id] = {}
            correlation_map[corr_id][event.get("cat", "")] = event
    # Second pass: extract collectives and compute events
    for event in events:
        if not isinstance(event, dict):
            continue
        cat = event.get("cat", "")
        name = event.get("name", "")
        ph = event.get("ph", "")

        if ph not in ["X", "B", "E"]:
            continue
        collective = _try_parse_collective(event, rank, correlation_map)
        if collective is not None:
            timeline.add_collective(collective)
            continue
        compute = _try_parse_compute(event)
        if compute is not None:
            timeline.add_compute_event(compute)
    # Infer device_id and pid from events
    for event in events:
        if isinstance(event, dict):
            args = event.get("args", {})
            if "device" in args:
                timeline.device_id = args["device"]
                break
    if events:
        first_event = events[0]
        if isinstance(first_event, dict):
            timeline.pid = first_event.get("pid")
    return timeline, None
def _infer_rank_from_filename(filename: str) -> int:
    """Try to infer rank from filename patterns.
    Common patterns:
    - trace_rank0.json, trace_rank_0.json
    - worker0/trace.json
    - rank0_trace.json
    """
    import re
    # Pattern: rank followed by number
    match = re.search(r"rank[_\-]?(\d+)", filename.lower())
    if match:
        return int(match.group(1))
    # Pattern: worker followed by number
    match = re.search(r"worker[_\-]?(\d+)", filename.lower())
    if match:
        return int(match.group(1))
    # Pattern: just a number before .json
    match = re.search(r"_(\d+)\.json", filename.lower())
    if match:
        return int(match.group(1))
    return 0
def _try_parse_collective(
    event: dict,
    rank: int,
    correlation_map: dict[int, dict],
) -> Optional[Collective]:
    """Try to parse an event as a collective operation.
    Identifies collectives by:
    1. Operator names: nccl:all_reduce, torch.distributed.all_reduce, etc.
    2. Kernel names: ncclDevKernel_*, rcclDevKernel_*
    """
    name = event.get("name", "")
    cat = event.get("cat", "")
    args = event.get("args", {})
    is_collective = False
    collective_type = CollectiveType.UNKNOWN
    # Pattern 1: nccl: or rccl: prefix in operator name
    if name.startswith("nccl:") or name.startswith("rccl:"):
        is_collective = True
        collective_type = CollectiveType.from_string(name)
    # Pattern 2: torch.distributed calls
    elif "distributed" in name.lower() and any(
        op in name.lower()
        for op in ["all_reduce", "all_gather", "reduce_scatter", "broadcast", "barrier"]
    ):
        is_collective = True
        collective_type = CollectiveType.from_string(name)
    # Pattern 3: NCCL/RCCL kernel names
    elif "ncclDevKernel" in name or "rcclDevKernel" in name:
        is_collective = True
        # Try to extract collective type from kernel name
        # e.g., ncclDevKernel_Generic -> look at surrounding context
        collective_type = _infer_collective_from_kernel(name, args)
    # Pattern 4: Category indicates NCCL
    elif cat in ["nccl", "rccl", "cuda_runtime"] and any(
        keyword in name.lower() for keyword in ["allreduce", "allgather", "reducescatter", "broadcast"]
    ):
        is_collective = True
        collective_type = CollectiveType.from_string(name)
    if not is_collective:
        return None
    ts = event.get("ts", 0)  # Timestamp in microseconds
    dur = event.get("dur", 0)  # Duration in microseconds
    start_time_ns = int(ts * 1000)
    end_time_ns = int((ts + dur) * 1000)
    message_size = args.get("message_size", args.get("bytes", args.get("size", 0)))
    if isinstance(message_size, str):
        try:
            message_size = int(message_size)
        except ValueError:
            message_size = 0
    count = args.get("count", args.get("nelem", 0))
    dtype_str = args.get("dtype", args.get("datatype", ""))
    datatype = DataType.from_string(dtype_str) if dtype_str else DataType.UNKNOWN
    if message_size == 0 and count > 0:
        message_size = count * datatype.size_bytes()
    pg_id = args.get("pg_id", args.get("comm_id", args.get("communicator", "")))
    pg_name = args.get("pg_name", args.get("process_group", None))
    process_group = ProcessGroup(comm_id=str(pg_id), name=pg_name) if pg_id else None
    corr_id = args.get("correlation", args.get("correlationId"))
    stack = args.get("stack", args.get("Python stack", args.get("call_stack")))
    stack_trace = "\n".join(stack) if isinstance(stack, list) else stack
    algorithm = args.get("algorithm", args.get("algo"))
    protocol = args.get("protocol", args.get("proto"))
    return Collective(
        collective_type=collective_type,
        rank=rank,
        start_time_ns=start_time_ns,
        end_time_ns=end_time_ns,
        message_size_bytes=message_size,
        count=count,
        datatype=datatype,
        process_group=process_group,
        correlation_id=corr_id,
        pytorch_op_name=name if not name.startswith("ncclDev") else None,
        pytorch_stack_trace=stack_trace,
        algorithm=algorithm,
        protocol=protocol,
        extra=dict(args) if args else {},
    )
def _infer_collective_from_kernel(name: str, args: dict) -> CollectiveType:
    """Infer collective type from kernel name and arguments."""
    name_lower = name.lower()
    if "allreduce" in name_lower:
        return CollectiveType.ALL_REDUCE
    if "allgather" in name_lower:
        return CollectiveType.ALL_GATHER
    if "reducescatter" in name_lower:
        return CollectiveType.REDUCE_SCATTER
    if "broadcast" in name_lower or "bcast" in name_lower:
        return CollectiveType.BROADCAST
    if "alltoall" in name_lower:
        return CollectiveType.ALL_TO_ALL
    op = args.get("op", args.get("operation", ""))
    if op:
        return CollectiveType.from_string(op)
    return CollectiveType.UNKNOWN
def _try_parse_compute(event: dict) -> Optional[ComputeEvent]:
    """Try to parse an event as a compute operation.
    Filters out NCCL/RCCL kernels to only include compute kernels.
    """
    name = event.get("name", "")
    cat = event.get("cat", "")
    args = event.get("args", {})
    # Only process kernel events
    if cat not in ["kernel", "cuda_kernel", "gpu_kernel", "hip_kernel"]:
        return None

    if any(keyword in name.lower() for keyword in ["nccl", "rccl", "allreduce", "allgather"]):
        return None
    ts = event.get("ts", 0)
    dur = event.get("dur", 0)
    if dur <= 0:
        return None
    start_time_ns = int(ts * 1000)
    end_time_ns = int((ts + dur) * 1000)
    grid = args.get("grid", args.get("gridDim"))
    block = args.get("block", args.get("blockDim"))
    grid_size = None
    if grid:
        if isinstance(grid, list) and len(grid) >= 3:
            grid_size = (grid[0], grid[1], grid[2])
        elif isinstance(grid, str):
            parts = grid.strip("()[]").split(",")
            if len(parts) >= 3:
                grid_size = (int(parts[0]), int(parts[1]), int(parts[2]))
    block_size = None
    if block:
        if isinstance(block, list) and len(block) >= 3:
            block_size = (block[0], block[1], block[2])
        elif isinstance(block, str):
            parts = block.strip("()[]").split(",")
            if len(parts) >= 3:
                block_size = (int(parts[0]), int(parts[1]), int(parts[2]))
    stream_id = args.get("stream", args.get("stream_id", 0))
    device_id = args.get("device", args.get("device_id", 0))
    shared_mem = args.get("shared_memory", args.get("sharedMemory", 0))
    registers = args.get("registers", args.get("registersPerThread", 0))
    corr_id = args.get("correlation", args.get("correlationId"))
    return ComputeEvent(
        name=name,
        start_time_ns=start_time_ns,
        end_time_ns=end_time_ns,
        stream_id=stream_id,
        device_id=device_id,
        grid_size=grid_size,
        block_size=block_size,
        shared_memory_bytes=shared_mem,
        registers_per_thread=registers,
        correlation_id=corr_id,
        extra=dict(args) if args else {},
    )
