#!/usr/bin/env python3
"""Read-only BAR0 probe for GA100 GPC floorsweep fuses (OPT_GPC_DISABLE / OPT_GPC_DEFECTIVE).

Never writes to BAR0. Requires root to mmap /sys/bus/pci/devices/<bdf>/resource0.
"""

import mmap
import os
import struct
import subprocess
import sys

REG_GPC_DISABLE = 0x00820350
REG_GPC_DEFECTIVE = 0x008205C4
PAGE = 0x1000


def find_ga100_cards(pattern="GA100"):
    """Return [(bdf, label)] for PCI devices whose lspci name matches pattern."""
    try:
        out = subprocess.check_output(["lspci", "-D"], text=True)
    except FileNotFoundError:
        sys.exit("lspci not found; install pciutils or pass BDFs another way")
    except subprocess.CalledProcessError as e:
        sys.exit(f"lspci failed: {e}")
    cards = []
    for line in out.splitlines():
        if pattern in line:
            bdf = line.split(None, 1)[0]
            cards.append((bdf, bdf.split(":")[1]))
    return cards


def read_reg(bdf, addr):
    path = f"/sys/bus/pci/devices/{bdf}/resource0"
    if not os.path.exists(path):
        raise OSError(f"no BAR0 (resource0 missing, device has no MMIO?)")
    if os.path.getsize(path) <= addr:
        raise OSError(f"BAR0 too small for offset 0x{addr:x}")
    fd = os.open(path, os.O_RDONLY | os.O_SYNC)
    try:
        base = addr - (addr % PAGE)
        off = addr - base
        m = mmap.mmap(fd, PAGE, mmap.MAP_SHARED, mmap.PROT_READ, offset=base)
        try:
            return struct.unpack_from("<I", m, off)[0]
        finally:
            m.close()
    finally:
        os.close(fd)


def bits_set(mask, width=8):
    return [i for i in range(width) if mask & (1 << i)]


def main():
    if os.geteuid() != 0:
        sys.exit("must run as root (mmaps BAR0 via /sys/.../resource0)")

    cards = find_ga100_cards()
    if not cards:
        sys.exit("no GA100 cards found (lspci name match)")
    print(f"Detected GA100 card(s): {', '.join(b for b, _ in cards)}\n")

    print(f"{'Card':<6} {'Disabled clusters':<20} {'Marked defective':<18} {'Disabled but not defective'}")
    for bdf, label in cards:
        try:
            disable = read_reg(bdf, REG_GPC_DISABLE)
            defective = read_reg(bdf, REG_GPC_DEFECTIVE)
        except OSError as e:
            print(f"{label:<6} ERROR reading {bdf}: {e}")
            continue

        disabled_bits = set(bits_set(disable))
        defective_bits = set(bits_set(defective))
        not_defective_bits = disabled_bits - defective_bits

        def fmt(bits):
            return ", ".join(str(b) for b in sorted(bits)) if bits else "none"

        print(f"{label:<6} {fmt(disabled_bits):<20} {fmt(defective_bits):<18} {fmt(not_defective_bits)}")
        print(f"       (raw: OPT_GPC_DISABLE=0x{disable:08x} OPT_GPC_DEFECTIVE=0x{defective:08x})")


if __name__ == "__main__":
    main()