"""ROCm Systems Profiler parsing.
Parses trace output from ROCm Systems Profiler (rocprof-sys) which can capture
RCCL operations when enabled with ROCPROFSYS_USE_RCCLP=ON.
The output can be in Perfetto format (.proto) or other formats.
Format Reference:
- https://rocm.docs.amd.com/projects/omnitrace/en/latest/
"""
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_rocsys_file(
    file_path: str | Path,
    rank: Optional[int] = None,
) -> tuple[RankTimeline | None, str | None]:
    """Parse a ROCm Systems Profiler output file.
    Supports:
    - JSON/JSONL trace files
    - Perfetto-compatible JSON
    - rocpd databases
    Args:
        file_path: Path to the trace 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}"
    suffix = path.suffix.lower()
    try:
        if suffix in [".rocpd", ".db"]:
            return _parse_rocpd_database(path, rank)
        elif suffix in [".json", ".perfetto-trace"]:
            return _parse_json_trace(path, rank)
        elif suffix == ".proto":
            return _parse_protobuf_trace(path, rank)
        else:
            # Try JSON first, then database
            try:
                return _parse_json_trace(path, rank)
            except Exception:
                return _parse_rocpd_database(path, rank)
    except Exception as e:
        return None, f"Failed to parse ROCm Systems Profiler output: {e}"
def parse_rocsys(
    data: dict | list,
    rank: Optional[int] = None,
    filename: Optional[str] = None,
) -> tuple[RankTimeline | None, str | None]:
    """Parse ROCm Systems Profiler trace data.
    Args:
        data: Parsed trace data (JSON dict or list of events)
        rank: Rank number. If None, will try to infer.
        filename: Original filename for rank inference
    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 isinstance(data, dict):
        events = data.get("traceEvents", data.get("trace_events", []))
        if not events and "data" in data:
            events = data["data"]
    else:
        events = data
    for event in events:
        if not isinstance(event, dict):
            continue
        # Try to parse as collective
        collective = _try_parse_collective(event, rank)
        if collective is not None:
            timeline.add_collective(collective)
            continue
        # Try to parse as compute
        compute = _try_parse_compute(event)
        if compute is not None:
            timeline.add_compute_event(compute)
    return timeline, None
def _parse_json_trace(
    path: Path,
    rank: Optional[int],
) -> tuple[RankTimeline | None, str | None]:
    """Parse JSON/Perfetto trace file."""
    with open(path, encoding="utf-8") as f:
        data = json.load(f)
    return parse_rocsys(data, rank, str(path))
def _parse_rocpd_database(
    path: Path,
    rank: Optional[int],
) -> tuple[RankTimeline | None, str | None]:
    """Parse rocpd SQLite database.
    
    rocpd databases contain tables for:
    - hip_api: HIP API calls
    - hsa_api: HSA API calls
    - roctracer_api: ROCtracer events
    - ext_op: External operations including RCCL
    """
    import sqlite3
    if rank is None:
        rank = _infer_rank_from_filename(str(path))
    timeline = RankTimeline(rank=rank)
    try:
        conn = sqlite3.connect(str(path))
        conn.row_factory = sqlite3.Row
        cursor = conn.execute(
            "SELECT name FROM sqlite_master WHERE type='table'"
        )
        tables = {row["name"] for row in cursor.fetchall()}
        if "ext_op" in tables:
            _parse_ext_op_table(conn, timeline)
        if "hip_api" in tables:
            _parse_hip_api_table(conn, timeline)
        if "kernel" in tables:
            _parse_kernel_table(conn, timeline)
        conn.close()
    except sqlite3.Error as e:
        return None, f"SQLite error: {e}"
    return timeline, None
def _parse_ext_op_table(conn: "sqlite3.Connection", timeline: RankTimeline) -> None:
    """Parse ext_op table for RCCL operations."""
    import sqlite3
    try:
        cursor = conn.execute("""
            SELECT BeginNs, EndNs, Name, pid, tid, args
            FROM ext_op
            WHERE Name LIKE '%rccl%' OR Name LIKE '%nccl%'
               OR Name LIKE '%allreduce%' OR Name LIKE '%allgather%'
               OR Name LIKE '%broadcast%' OR Name LIKE '%reducescatter%'
        """)
        for row in cursor:
            row_dict = dict(row)
            collective = _create_collective_from_rocpd(row_dict, timeline.rank)
            if collective is not None:
                timeline.add_collective(collective)
    except sqlite3.Error:
        pass
def _parse_hip_api_table(conn: "sqlite3.Connection", timeline: RankTimeline) -> None:
    """Parse hip_api table for kernel-related API calls."""
    import sqlite3
    try:
        cursor = conn.execute("""
            SELECT BeginNs, EndNs, Name, args, correlation_id
            FROM hip_api
            WHERE Name LIKE '%Launch%'
            LIMIT 10000
        """)
        # We primarily care about this for correlation, not direct compute events
        # The actual kernel timing is in the kernel table
    except sqlite3.Error:
        pass
def _parse_kernel_table(conn: "sqlite3.Connection", timeline: RankTimeline) -> None:
    """Parse kernel table for GPU kernel executions."""
    import sqlite3
    try:
        cursor = conn.execute("""
            SELECT BeginNs, EndNs, Name, DeviceId, QueueId, correlation_id
            FROM kernel
            WHERE Name NOT LIKE '%rccl%' AND Name NOT LIKE '%nccl%'
            LIMIT 10000
        """)
        for row in cursor:
            row_dict = dict(row)
            compute = _create_compute_from_rocpd(row_dict)
            if compute is not None:
                timeline.add_compute_event(compute)
    except sqlite3.Error:
        pass
def _parse_protobuf_trace(
    path: Path,
    rank: Optional[int],
) -> tuple[RankTimeline | None, str | None]:
    """Parse Perfetto protobuf trace.
    
    For now, we convert to JSON using perfetto trace_processor if available,
    otherwise return an error suggesting the user export to JSON.
    """
    # Try to use our Perfetto trace processor if available
    try:
        from wafer_core.lib.perfetto import PerfettoTool
        # This would require setting up the tool properly
        # For now, return an error with instructions
        return None, (
            "Protobuf trace format requires conversion. "
            "Please export the trace to JSON format using: "
            "trace_processor --json <file.proto>"
        )
    except ImportError:
        return None, (
            "Protobuf trace format not supported. "
            "Please export the trace to JSON format."
        )
def _try_parse_collective(event: dict, rank: int) -> Optional[Collective]:
    """Try to parse an event as a collective operation."""
    name = event.get("name", event.get("Name", ""))
    cat = event.get("cat", event.get("category", "")).lower()
    name_lower = name.lower()
    is_collective = any(
        pattern in name_lower
        for pattern in ["rccl", "nccl", "allreduce", "allgather", "reducescatter", "broadcast", "barrier"]
    ) or cat in ["rccl", "nccl", "collective"]
    if not is_collective:
        return None
    ts = event.get("ts", event.get("BeginNs", event.get("startNs", 0)))
    dur = event.get("dur", 0)
    end = event.get("EndNs", event.get("endNs"))
    if end is None:
        end = ts + dur
    if ts < 1e12:  # Likely microseconds
        ts = int(ts * 1000)
        end = int(end * 1000) if end else ts
    collective_type = CollectiveType.from_string(name)
    args = event.get("args", {})
    message_size = args.get("bytes", args.get("message_size", args.get("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()
    return Collective(
        collective_type=collective_type,
        rank=rank,
        start_time_ns=int(ts),
        end_time_ns=int(end),
        message_size_bytes=message_size,
        count=count,
        datatype=datatype,
        extra=dict(event),
    )
def _try_parse_compute(event: dict) -> Optional[ComputeEvent]:
    """Try to parse an event as a compute operation."""
    name = event.get("name", event.get("Name", ""))
    cat = event.get("cat", event.get("category", "")).lower()

    if cat not in ["kernel", "hip", "gpu", "compute"]:
        return None

    if any(pattern in name.lower() for pattern in ["rccl", "nccl"]):
        return None
    ts = event.get("ts", event.get("BeginNs", 0))
    dur = event.get("dur", 0)
    end = event.get("EndNs")
    if end is None:
        end = ts + dur
    if ts < 1e12:  # Convert from microseconds
        ts = int(ts * 1000)
        end = int(end * 1000) if end else ts
    if end <= ts:
        return None
    args = event.get("args", {})
    stream_id = args.get("stream", args.get("QueueId", 0)) or 0
    device_id = args.get("device", args.get("DeviceId", 0)) or 0
    corr_id = args.get("correlation_id", event.get("correlation_id"))
    return ComputeEvent(
        name=name,
        start_time_ns=int(ts),
        end_time_ns=int(end),
        stream_id=stream_id,
        device_id=device_id,
        correlation_id=int(corr_id) if corr_id else None,
        extra=dict(event),
    )
def _create_collective_from_rocpd(row: dict, rank: int) -> Optional[Collective]:
    """Create a Collective from a rocpd database row."""
    start_ns = row.get("BeginNs")
    end_ns = row.get("EndNs")
    name = row.get("Name", "")
    if start_ns is None or end_ns is None:
        return None
    collective_type = CollectiveType.from_string(name)
    args_str = row.get("args", "{}")
    try:
        args = json.loads(args_str) if isinstance(args_str, str) else {}
    except json.JSONDecodeError:
        args = {}
    message_size = args.get("bytes", args.get("size", 0))
    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,
        extra=row,
    )
def _create_compute_from_rocpd(row: dict) -> Optional[ComputeEvent]:
    """Create a ComputeEvent from a rocpd database row."""
    start_ns = row.get("BeginNs")
    end_ns = row.get("EndNs")
    name = row.get("Name", "")
    if start_ns is None or end_ns is None or not name:
        return None
    device_id = row.get("DeviceId", 0) or 0
    stream_id = row.get("QueueId", 0) or 0
    corr_id = row.get("correlation_id")
    return ComputeEvent(
        name=name,
        start_time_ns=int(start_ns),
        end_time_ns=int(end_ns),
        device_id=device_id,
        stream_id=stream_id,
        correlation_id=int(corr_id) if corr_id else None,
        extra=dict(row),
    )
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+)\.(rocpd|db|json|proto)", filename.lower())
    if match:
        return int(match.group(1))
    return 0
