"""Nsight Systems SQLite parsing.
Parses .nsys-rep files which are SQLite databases containing GPU timeline,
CUDA events, and NCCL operations when profiled with:
    nsys profile -t cuda,nccl --stats=true <command>
The SQLite database contains tables like:
- CUDA_KERNEL_CALLS
- NCCL_EVENTS (when NCCL tracing enabled)
- NVTX_EVENTS
- CUPTI_ACTIVITY_KIND_*
Format Reference:
- https://docs.nvidia.com/nsight-systems/UserGuide/index.html
"""
import sqlite3
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_nsys_file(
    file_path: str | Path,
    rank: Optional[int] = None,
) -> tuple[RankTimeline | None, str | None]:
    """Parse an Nsight Systems report file.
    Args:
        file_path: Path to .nsys-rep or .sqlite 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:
        return parse_nsys(str(path), rank)
    except Exception as e:
        return None, f"Failed to parse Nsight Systems report: {e}"
def parse_nsys(
    file_path: str,
    rank: Optional[int] = None,
) -> tuple[RankTimeline | None, str | None]:
    """Parse Nsight Systems SQLite database.
    Args:
        file_path: Path to the SQLite database
        rank: Rank number. If None, will try to infer.
    Returns:
        Tuple of (RankTimeline, error). One will be None.
    """
    if rank is None:
        rank = _infer_rank_from_filename(file_path)
    timeline = RankTimeline(rank=rank)
    try:
        conn = sqlite3.connect(file_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 "NCCL_EVENTS" in tables:
            _parse_nccl_events(conn, timeline)
        kernel_tables = [t for t in tables if "KERNEL" in t.upper()]
        for table in kernel_tables:
            _parse_kernel_table(conn, table, timeline)
        if "NVTX_EVENTS" in tables:
            _parse_nvtx_events(conn, timeline)
        # Try to get process info
        if "TARGET_INFO_SYSTEM_ENV" in tables:
            _extract_process_info(conn, timeline)
        conn.close()
    except sqlite3.Error as e:
        return None, f"SQLite error: {e}"
    return timeline, None
def _parse_nccl_events(conn: sqlite3.Connection, timeline: RankTimeline) -> None:
    """Parse NCCL_EVENTS table for collective operations."""
    try:
        # Try different possible column schemas
        cursor = conn.execute("PRAGMA table_info(NCCL_EVENTS)")
        columns = {row["name"].lower() for row in cursor.fetchall()}
        select_cols = []
        if "start" in columns:
            select_cols.append("start")
        if "end" in columns:
            select_cols.append("end")
        if "startnstamp" in columns:
            select_cols.append("startNStamp as start")
        if "endnstamp" in columns:
            select_cols.append("endNStamp as end")
        if "collectivename" in columns:
            select_cols.append("collectiveName")
        if "optype" in columns:
            select_cols.append("opType")
        if "sendbytes" in columns:
            select_cols.append("sendBytes")
        if "recvbytes" in columns:
            select_cols.append("recvBytes")
        if "commid" in columns:
            select_cols.append("commId")
        if "rank" in columns:
            select_cols.append("rank")
        if "globalrank" in columns:
            select_cols.append("globalRank")
        if not select_cols:
            select_cols = ["*"]
        query = f"SELECT {', '.join(select_cols)} FROM NCCL_EVENTS"
        cursor = conn.execute(query)
        for row in cursor:
            row_dict = dict(row)
            collective = _create_collective_from_nsys_row(row_dict, timeline.rank)
            if collective is not None:
                timeline.add_collective(collective)
    except sqlite3.Error:
        # Table might not exist or have different schema
        pass
def _parse_kernel_table(
    conn: sqlite3.Connection,
    table: str,
    timeline: RankTimeline,
) -> None:
    """Parse a CUDA kernel table."""
    try:
        cursor = conn.execute(f"SELECT * FROM {table} LIMIT 1000")
        for row in cursor:
            row_dict = dict(row)
            
            name = str(row_dict.get("shortName", row_dict.get("demangledName", row_dict.get("name", ""))))
            if any(pattern in name.lower() for pattern in ["nccl", "rccl"]):
                continue
            compute = _create_compute_from_nsys_row(row_dict)
            if compute is not None:
                timeline.add_compute_event(compute)
    except sqlite3.Error:
        pass
def _parse_nvtx_events(conn: sqlite3.Connection, timeline: RankTimeline) -> None:
    """Parse NVTX events for additional collective context.
    
    Some frameworks use NVTX markers to annotate collective operations.
    """
    try:
        cursor = conn.execute("""
            SELECT text, start, end, globalTid
            FROM NVTX_EVENTS
            WHERE text LIKE '%allreduce%' 
               OR text LIKE '%all_reduce%'
               OR text LIKE '%allgather%'
               OR text LIKE '%broadcast%'
               OR text LIKE '%reducescatter%'
            LIMIT 1000
        """)
        for row in cursor:
            row_dict = dict(row)
            text = row_dict.get("text", "")
            start = row_dict.get("start")
            end = row_dict.get("end")
            if start is None or end is None:
                continue
            collective_type = CollectiveType.from_string(text)
            if collective_type == CollectiveType.UNKNOWN:
                continue
            collective = Collective(
                collective_type=collective_type,
                rank=timeline.rank,
                start_time_ns=int(start),
                end_time_ns=int(end),
                extra={"source": "nvtx", "text": text},
            )
            timeline.add_collective(collective)
    except sqlite3.Error:
        pass
def _extract_process_info(conn: sqlite3.Connection, timeline: RankTimeline) -> None:
    """Extract process information from the database."""
    try:
        cursor = conn.execute("""
            SELECT name, value FROM TARGET_INFO_SYSTEM_ENV
            WHERE name IN ('RANK', 'LOCAL_RANK', 'WORLD_SIZE', 'HOSTNAME')
        """)
        for row in cursor:
            name = row["name"]
            value = row["value"]
            if name == "RANK" and timeline.rank == 0:
                try:
                    timeline.rank = int(value)
                except ValueError:
                    pass
            elif name == "HOSTNAME":
                timeline.hostname = value
    except sqlite3.Error:
        pass
def _create_collective_from_nsys_row(
    row: dict,
    rank: int,
) -> Optional[Collective]:
    """Create a Collective from an NSYS NCCL_EVENTS row."""
    start_ns = row.get("start", row.get("startNStamp"))
    end_ns = row.get("end", row.get("endNStamp"))
    if start_ns is None or end_ns is None:
        return None
    op_name = row.get("collectiveName", row.get("opType", ""))
    collective_type = CollectiveType.from_string(op_name)
    send_bytes = row.get("sendBytes", 0) or 0
    recv_bytes = row.get("recvBytes", 0) or 0
    message_size = max(send_bytes, recv_bytes)
    comm_id = row.get("commId", row.get("comm_id", ""))
    comm_rank = row.get("rank", row.get("globalRank", rank))
    process_group = ProcessGroup(comm_id=str(comm_id)) if comm_id else None
    return Collective(
        collective_type=collective_type,
        rank=comm_rank if comm_rank is not None else rank,
        start_time_ns=int(start_ns),
        end_time_ns=int(end_ns),
        message_size_bytes=message_size,
        process_group=process_group,
        extra=dict(row),
    )
def _create_compute_from_nsys_row(row: dict) -> Optional[ComputeEvent]:
    """Create a ComputeEvent from an NSYS kernel row."""
    name = row.get("shortName", row.get("demangledName", row.get("name", "")))
    if not name:
        return None
    start_ns = row.get("start", row.get("startNStamp"))
    end_ns = row.get("end", row.get("endNStamp"))
    if start_ns is None or end_ns is None:
        return None
    stream_id = row.get("streamId", row.get("stream", 0)) or 0
    device_id = row.get("deviceId", row.get("device", 0)) or 0
    grid_x = row.get("gridX", 1)
    grid_y = row.get("gridY", 1)
    grid_z = row.get("gridZ", 1)
    grid_size = (grid_x, grid_y, grid_z) if grid_x else None
    block_x = row.get("blockX", 1)
    block_y = row.get("blockY", 1)
    block_z = row.get("blockZ", 1)
    block_size = (block_x, block_y, block_z) if block_x else None
    shared_mem = row.get("sharedMemory", row.get("staticSharedMemory", 0)) or 0
    registers = row.get("registersPerThread", row.get("registers", 0)) or 0
    corr_id = row.get("correlationId")
    return ComputeEvent(
        name=str(name),
        start_time_ns=int(start_ns),
        end_time_ns=int(end_ns),
        stream_id=int(stream_id),
        device_id=int(device_id),
        grid_size=grid_size,
        block_size=block_size,
        shared_memory_bytes=int(shared_mem),
        registers_per_thread=int(registers),
        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+)\.(nsys-rep|sqlite)", filename.lower())
    if match:
        return int(match.group(1))
    return 0
