"""rocprofv3 RCCL CSV/JSON parsing.
Parses trace output from rocprofv3 when profiling RCCL operations with:
    rocprofv3 --rccl-trace --output-format csv <command>
    rocprofv3 --rccl-trace --output-format json <command>
CSV Format columns:
    Domain, Function, Process_Id, Thread_Id, Correlation_Id, Start_Timestamp, End_Timestamp
Format Reference:
- https://rocm.docs.amd.com/projects/rocprofiler-sdk/en/develop/how-to/using-rocprofv3.html
"""
import csv
import json
from pathlib import Path
from typing import Optional
from wafer_core.lib.distributed_traces.models.collective import (
    Collective,
    CollectiveType,
    DataType,
)
from wafer_core.lib.distributed_traces.models.rank_timeline import ComputeEvent, RankTimeline
def parse_rocprofv3_file(
    file_path: str | Path,
    rank: Optional[int] = None,
) -> tuple[RankTimeline | None, str | None]:
    """Parse a rocprofv3 output file (CSV or JSON).
    Args:
        file_path: Path to the CSV or JSON 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:
        suffix = path.suffix.lower()
        if suffix == ".csv":
            with open(path, encoding="utf-8") as f:
                reader = csv.DictReader(f)
                rows = list(reader)
            return parse_rocprofv3(rows, rank, str(path), is_csv=True)
        elif suffix == ".json":
            with open(path, encoding="utf-8") as f:
                data = json.load(f)
            return parse_rocprofv3(data, rank, str(path), is_csv=False)
        else:
            # Try CSV first, then JSON
            try:
                with open(path, encoding="utf-8") as f:
                    reader = csv.DictReader(f)
                    rows = list(reader)
                return parse_rocprofv3(rows, rank, str(path), is_csv=True)
            except Exception:
                with open(path, encoding="utf-8") as f:
                    data = json.load(f)
                return parse_rocprofv3(data, rank, str(path), is_csv=False)
    except Exception as e:
        return None, f"Failed to parse rocprofv3 output: {e}"
def parse_rocprofv3(
    data: list[dict] | dict,
    rank: Optional[int] = None,
    filename: Optional[str] = None,
    is_csv: bool = True,
) -> tuple[RankTimeline | None, str | None]:
    """Parse rocprofv3 output data.
    Args:
        data: List of row dicts (CSV) or JSON data
        rank: Rank number. If None, will try to infer.
        filename: Original filename for rank inference
        is_csv: True if data is from CSV, False if from JSON
    Returns:
        Tuple of (RankTimeline, error). One will be None.
    """
    if rank is None:
        rank = _infer_rank_from_filename(filename) if filename else 0
    timeline = RankTimeline(rank=rank)
    if is_csv:
        _parse_csv_rows(data, timeline)
    else:
        # JSON format might be a list or have a data key
        rows = data.get("data", data) if isinstance(data, dict) else data
        if isinstance(rows, list):
            _parse_json_rows(rows, timeline)
    return timeline, None
def _parse_csv_rows(rows: list[dict], timeline: RankTimeline) -> None:
    """Parse CSV rows into timeline events."""

    correlation_map: dict[int, dict] = {}
    for row in rows:
        domain = row.get("Domain", row.get("domain", "")).lower()
        function = row.get("Function", row.get("function", ""))
        
        # Normalize column names (rocprofv3 uses various conventions)
        start_ts = _get_timestamp(row, ["Start_Timestamp", "start_timestamp", "StartNs", "start"])
        end_ts = _get_timestamp(row, ["End_Timestamp", "end_timestamp", "EndNs", "end"])
        corr_id = _get_int(row, ["Correlation_Id", "correlation_id", "CorrelationId"])
        process_id = _get_int(row, ["Process_Id", "process_id", "ProcessId", "pid"])
        thread_id = _get_int(row, ["Thread_Id", "thread_id", "ThreadId", "tid"])
        if start_ts is None or end_ts is None:
            continue

        if corr_id is not None:
            correlation_map[corr_id] = row
        if _is_rccl_operation(domain, function):
            collective = _create_collective_from_row(
                function, start_ts, end_ts, timeline.rank, row
            )
            if collective is not None:
                timeline.add_collective(collective)
            continue
        if domain in ["hip", "kernel", "gpu"] and not _is_rccl_operation(domain, function):
            compute = _create_compute_from_row(
                function, start_ts, end_ts, row
            )
            if compute is not None:
                timeline.add_compute_event(compute)
def _parse_json_rows(rows: list[dict], timeline: RankTimeline) -> None:
    """Parse JSON format rows into timeline events."""
    for row in rows:
        # JSON format might use different field names
        function = row.get("name", row.get("function", row.get("kernel_name", "")))
        kind = row.get("kind", row.get("domain", "")).lower()
        start_ts = row.get("start", row.get("begin", row.get("startNs")))
        end_ts = row.get("end", row.get("endNs"))
        duration = row.get("dur", row.get("duration"))
        if start_ts is None:
            continue
        if end_ts is None and duration is not None:
            end_ts = start_ts + duration
        elif end_ts is None:
            continue
        if _is_rccl_operation(kind, function):
            collective = _create_collective_from_row(
                function, int(start_ts), int(end_ts), timeline.rank, row
            )
            if collective is not None:
                timeline.add_collective(collective)
            continue
        if kind in ["kernel", "hip_kernel", "compute"]:
            compute = _create_compute_from_row(
                function, int(start_ts), int(end_ts), row
            )
            if compute is not None:
                timeline.add_compute_event(compute)
def _is_rccl_operation(domain: str, function: str) -> bool:
    """Check if this is an RCCL collective operation."""
    function_lower = function.lower()
    domain_lower = domain.lower()
    if domain_lower in ["rccl", "nccl", "collective"]:
        return True
    rccl_patterns = [
        "rccl", "nccl",
        "allreduce", "all_reduce",
        "allgather", "all_gather",
        "reducescatter", "reduce_scatter",
        "broadcast", "bcast",
        "alltoall", "all_to_all",
        "barrier",
    ]
    return any(pattern in function_lower for pattern in rccl_patterns)
def _create_collective_from_row(
    function: str,
    start_ns: int,
    end_ns: int,
    rank: int,
    row: dict,
) -> Optional[Collective]:
    """Create a Collective from a row."""
    collective_type = CollectiveType.from_string(function)
    count = _get_int(row, ["count", "nelem", "elements"])
    dtype_str = row.get("dtype", row.get("datatype", ""))
    datatype = DataType.from_string(dtype_str) if dtype_str else DataType.UNKNOWN
    message_size = _get_int(row, ["bytes", "message_size", "size"]) or 0
    if message_size == 0 and count and count > 0:
        message_size = count * datatype.size_bytes()
    corr_id = _get_int(row, ["Correlation_Id", "correlation_id", "CorrelationId"])
    return Collective(
        collective_type=collective_type,
        rank=rank,
        start_time_ns=start_ns,
        end_time_ns=end_ns,
        message_size_bytes=message_size,
        count=count or 0,
        datatype=datatype,
        correlation_id=corr_id,
        extra=dict(row),
    )
def _create_compute_from_row(
    function: str,
    start_ns: int,
    end_ns: int,
    row: dict,
) -> Optional[ComputeEvent]:
    """Create a ComputeEvent from a row."""
    stream_id = _get_int(row, ["stream", "stream_id", "StreamId"]) or 0
    device_id = _get_int(row, ["device", "device_id", "DeviceId"]) or 0
    corr_id = _get_int(row, ["Correlation_Id", "correlation_id", "CorrelationId"])
    grid_size = None
    block_size = None
    grid_x = _get_int(row, ["grid_x", "gridX"])
    grid_y = _get_int(row, ["grid_y", "gridY"])
    grid_z = _get_int(row, ["grid_z", "gridZ"])
    if grid_x is not None:
        grid_size = (grid_x, grid_y or 1, grid_z or 1)
    block_x = _get_int(row, ["block_x", "blockX"])
    block_y = _get_int(row, ["block_y", "blockY"])
    block_z = _get_int(row, ["block_z", "blockZ"])
    if block_x is not None:
        block_size = (block_x, block_y or 1, block_z or 1)
    return ComputeEvent(
        name=function,
        start_time_ns=start_ns,
        end_time_ns=end_ns,
        stream_id=stream_id,
        device_id=device_id,
        grid_size=grid_size,
        block_size=block_size,
        correlation_id=corr_id,
        extra=dict(row),
    )
def _get_timestamp(row: dict, keys: list[str]) -> Optional[int]:
    """Get timestamp value from row, trying multiple keys."""
    for key in keys:
        if key in row:
            val = row[key]
            try:
                return int(val)
            except (ValueError, TypeError):
                continue
    return None
def _get_int(row: dict, keys: list[str]) -> Optional[int]:
    """Get integer value from row, trying multiple keys."""
    for key in keys:
        if key in row:
            val = row[key]
            try:
                return int(val)
            except (ValueError, TypeError):
                continue
    return 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+)\.(csv|json)", filename.lower())
    if match:
        return int(match.group(1))
    return 0
