"""NVLink/PCIe topology detection via nvidia-smi topo -m."""

from __future__ import annotations

import logging
import subprocess
from typing import Optional

logger = logging.getLogger("radix_edge.observer")


def detect_topology() -> Optional[dict[int, list[int]]]:
    """Detect NVLink topology. Returns {gpu_index: [peer_indices]} or None if unavailable.

    Parses `nvidia-smi topo -m` output which looks like:
        GPU0  GPU1  GPU2  GPU3
    GPU0   X    NV4   NV4   NV4
    GPU1  NV4    X    NV4   NV4
    GPU2  NV4   NV4    X    NV4
    GPU3  NV4   NV4   NV4    X

    NVx = NVLink connection, SYS/PHB/PIX/PXB = PCIe connections.
    """
    try:
        result = subprocess.run(
            ["nvidia-smi", "topo", "-m"],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if result.returncode != 0:
            return None
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return None

    return parse_topology_output(result.stdout)


def parse_topology_output(output: str) -> Optional[dict[int, list[int]]]:
    """Parse nvidia-smi topo -m output into NVLink peer map."""
    lines = output.strip().split("\n")
    if len(lines) < 2:
        return None

    # Find the header line with GPU labels.
    # The header is the line that contains GPU0, GPU1, etc. but does NOT start with
    # "GPUn  X" (which would be a data row). Real nvidia-smi output may be tab or
    # space separated.
    header_idx = None
    for i, line in enumerate(lines):
        stripped = line.strip()
        if not stripped:
            continue
        cols = stripped.split()
        gpu_cols_in_line = [c for c in cols if c.startswith("GPU")]
        # Header line: has multiple GPU labels and the first non-GPU-label token
        # is NOT a connection type (X, NV*, PHB, SYS, PIX, PXB, N/A)
        if len(gpu_cols_in_line) >= 2:
            non_gpu = [c for c in cols if not c.startswith("GPU")]
            connection_types = {"X", "PHB", "SYS", "PIX", "PXB", "N/A", "SOC", "NODE"}
            is_data_row = any(c in connection_types or c.startswith("NV") for c in non_gpu[:4])
            if not is_data_row:
                header_idx = i
                break

    if header_idx is None:
        return None

    # Parse header to get GPU count
    header = lines[header_idx]
    gpu_cols = [col for col in header.split() if col.startswith("GPU")]
    n_gpus = len(gpu_cols)

    if n_gpus == 0:
        return None

    topology: dict[int, list[int]] = {i: [] for i in range(n_gpus)}

    # Parse connection matrix
    for line in lines[header_idx + 1:]:
        parts = line.split()
        if not parts or not parts[0].startswith("GPU"):
            continue

        try:
            src_idx = int(parts[0][3:])
        except ValueError:
            continue

        # Connection types start after the row label
        for dst_idx, conn in enumerate(parts[1:1 + n_gpus]):
            if dst_idx == src_idx:
                continue
            # NVx means NVLink connection (NV1, NV2, NV4, NV6, NV12, etc.)
            if conn.startswith("NV") and conn != "N/A":
                if dst_idx not in topology.get(src_idx, []):
                    topology.setdefault(src_idx, []).append(dst_idx)

    # Return None if no NVLink connections found
    has_nvlink = any(len(peers) > 0 for peers in topology.values())
    return topology if has_nvlink else None
