#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# SPDX-License-Identifier: MulanPSL-2.0
"""
gpudiag - GPU Diagnostic Tool

Detects the following GPU anomalies:
  1.  Xid errors (NVIDIA specific, from dmesg / log files)
  2.  GPU driver not loaded
  3.  SMI tool not found
  4.  SMI tool hung or returned error
  5.  Card drop (lspci vs smi count mismatch)
  6.  Z/D state processes holding GPU resources
  7.  dmesg driver error patterns (NVIDIA / AMD / Intel)
  8.  PCIe link degradation
  9.  ECC memory errors
  10. High temperature / thermal throttling
  11. Power near limit
  12. Memory utilization near full
  13. GPU utilization with no registered process (process leak)
  14. NVLink anomaly
  15. GPU reset / hang in dmesg
  16. PCIe AER errors
"""

import os
import re
import sys
import json
import datetime
import subprocess
import shutil
import argparse

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
SMI_TIMEOUT  = 10       # seconds to wait for smi commands
DMESG_LINES  = 2000     # tail lines from dmesg to inspect
LOG_FILES    = [
    "/var/log/messages",
    "/var/log/syslog",
    "/var/log/kern.log",
]


# ---------------------------------------------------------------------------
# GPU architectures
# ---------------------------------------------------------------------------
ARCH_NVIDIA  = "nvidia"
ARCH_AMD     = "amd"
ARCH_INTEL   = "intel"
ARCH_PPU     = "ppu"       # Alibaba PPU
ARCH_NPU     = "npu"       # Ascend NPU
ARCH_UNKNOWN = "unknown"

ALL_ARCHS = [ARCH_NVIDIA, ARCH_AMD, ARCH_INTEL, ARCH_PPU, ARCH_NPU]

# ---------------------------------------------------------------------------
# Per-vendor lspci detection rules
# Each entry: (arch, primary_grep_re, vendor_grep_re_or_None)
#   - primary_grep_re : first-pass filter to narrow lspci to GPU-like devices
#   - vendor_grep_re  : second-pass filter to identify the specific vendor
#                         None means primary_grep_re alone is sufficient
# ---------------------------------------------------------------------------
LSPCI_VENDOR_RULES = [
    # NVIDIA: 3D/VGA compatible controller + nvidia keyword
    (ARCH_NVIDIA, r"(?:3D|VGA compatible) controller", r"(?i)nvidia"),
    # Alibaba PPU: 3D/VGA compatible controller + alibaba or vendor:device 1ded:6001
    (ARCH_PPU,    r"(?:3D|VGA compatible) controller", r"(?i)(?:alibaba|1ded:6001)"),
    # AMD: "Processing accelerators" line with AMD name
    (ARCH_AMD,    r"Processing accelerators.*Advanced Micro Devices", None),
    # Ascend NPU: Device 6000 + Non-VGA
    (ARCH_NPU,    r"Device 6000", r"Non-VGA"),
    # Intel: VGA/Display controller + intel keyword (checked last — usually iGPU)
    (ARCH_INTEL,  r"(?:VGA|Display) controller",          r"(?i)intel"),
]

# SMI tool preference per architecture
# Only nvidia-smi and ppu-smi are actually supported in our environment:
#   - NVIDIA  (PCI vendor 0x10de) -> nvidia-smi
#   - Alibaba PPU (PCI vendor 0x1ded) -> ppu-smi
# Other architectures (AMD, Intel, NPU) do NOT have SMI tool support;
# SMI-dependent checks will be reported as "not supported" for them.
# Note: gpudiag runs on the node directly (NOT inside a container),
# so we do NOT prepend `nsenter -t 1 -m` like pouch_pid does.
SMI_PREFERRED = {
    ARCH_NVIDIA:  ["nvidia-smi"],
    ARCH_AMD:     [],   # rocm-smi not supported yet
    ARCH_PPU:     ["ppu-smi"],
    ARCH_INTEL:   [],   # not supported
    ARCH_NPU:     [],   # not supported
    ARCH_UNKNOWN: ["nvidia-smi", "ppu-smi"],
}

# Driver kernel module names per architecture
DRIVER_MODULES = {
    ARCH_NVIDIA:  ["nvidia"],
    ARCH_AMD:     ["amdgpu", "radeon"],
    ARCH_INTEL:   ["i915", "xe"],
    ARCH_PPU:     ["ppu"],
    ARCH_NPU:     ["hisilicon", "drv_hisi_hcc"],
    ARCH_UNKNOWN: ["nvidia", "amdgpu", "radeon", "i915"],
}

# ---------------------------------------------------------------------------
# dmesg patterns: (regex, severity, human description)
# ---------------------------------------------------------------------------
DMESG_PATTERNS = [
    # ── NVIDIA ─────────────────────────────────────────────────────────────
    (r"NVRM:.*GPU.*fallen off the bus",      "ERROR", "GPU fallen off PCIe bus"),
    (r"NVRM:.*GPU.*lost",                    "ERROR", "NVIDIA GPU lost on PCIe bus"),
    (r"NVRM:.*GPU-[0-9a-f].*not present",    "ERROR", "NVIDIA GPU device not present"),
    (r"(?i)pci.*(?:not present|device.*inaccessible|bus.*error.*ff)",
                                             "ERROR", "PCIe GPU device inaccessible/not present"),
    (r"NVRM:.*RmInitAdapter failed",         "ERROR", "NVIDIA adapter init failed"),
    (r"NVRM:.*error.*(?:timeout|reset)",     "ERROR", "NVIDIA driver timeout/reset error"),
    (r"NVRM:.*failed to allocate",           "ERROR", "NVIDIA driver resource allocation fail"),
    (r"NVRM:.*GPU.*power.*fault",            "ERROR", "NVIDIA GPU power fault"),
    (r"Xid.*\b(?:79|80|92|94|95|119)\b",     "ERROR", "Fatal Xid error (engine failure/hw fault)"),
    (r"NVRM:.*GPU Board.*not supported",     "ERROR", "Unsupported GPU board"),
    (r"NVRM:.*warning.*ECC",                 "WARN",  "NVIDIA ECC warning"),
    (r"NVRM:.*Nvlink.*(?:error|fault|fail)", "WARN",  "NVLink error/fault"),
    (r"NVRM:.*thermal.*throttl",             "WARN",  "NVIDIA thermal throttling"),
    (r"NVRM:.*power.*capping",               "WARN",  "NVIDIA power capping active"),
    # ── AMD ────────────────────────────────────────────────────────────────
    (r"amdgpu.*fault.*addr",                 "ERROR", "AMD GPU page fault"),
    (r"amdgpu.*ring.*timeout",               "ERROR", "AMD GPU ring/command timeout"),
    (r"amdgpu.*hw_init.*failed",             "ERROR", "AMD GPU hardware init failed"),
    (r"amdgpu.*GPU reset",                   "ERROR", "AMD GPU reset triggered"),
    (r"amdgpu.*ras.*uncorrectable",          "ERROR", "AMD GPU uncorrectable RAS error"),
    (r"amdgpu.*hung",                        "ERROR", "AMD GPU hang detected"),
    (r"amdgpu.*ras.*correctable",            "WARN",  "AMD GPU correctable RAS error"),
    (r"amdgpu.*thermal.*throttl",            "WARN",  "AMD GPU thermal throttling"),
    # ── Intel ──────────────────────────────────────────────────────────────
    (r"i915.*reset.*failed",                 "ERROR", "Intel GPU reset failed"),
    (r"i915.*gpu.*hang",                     "ERROR", "Intel GPU hang detected"),
    (r"i915.*engine.*reset",                 "WARN",  "Intel GPU engine reset"),
    (r"xe.*gt.*reset",                       "WARN",  "Intel Xe GPU GT reset"),
    # ── Generic PCIe ───────────────────────────────────────────────────────
    (r"pcieport.*AER.*Uncorrected",          "ERROR", "PCIe AER uncorrected error"),
    (r"pcieport.*AER.*Corrected",            "WARN",  "PCIe AER corrected error"),
    (r"pci.*error.*status",                  "WARN",  "PCI device error status"),
]

# ---------------------------------------------------------------------------
# Result accumulation
# ---------------------------------------------------------------------------
# Each entry: dict with keys: check, severity, message [, suggestion]
# This tool only reports raw observations — it does NOT classify the host
# as healthy/unhealthy. "severity" is just a description of the observation
# itself (OK / WARN / ERROR), no aggregate health verdict is emitted.
_results = []
# Diagnostic process log lines (appended by _log() and _report())
_logs = []


def _log(message):
    """Append a plain informational line to the process log."""
    _logs.append(message)


def _report(check, severity, message, suggestion=None):
    """Record a diagnostic result and append a log line."""
    entry = {
        "check":    check,
        "severity": severity,
        "message":  message,
    }
    if suggestion is not None:
        entry["suggestion"] = suggestion
    _results.append(entry)
    _logs.append(f"[{severity:<5}] {check}: {message}")

# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------
def _run(cmd, timeout=SMI_TIMEOUT):
    """
    Run a shell command.
    Returns (stdout, stderr, returncode).
    On timeout  -> stdout=None, returncode=-2
    On OSError  -> stdout=None, returncode=-3
    """
    try:
        proc = subprocess.run(
            cmd, shell=True,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            timeout=timeout, universal_newlines=True,
        )
        return proc.stdout, proc.stderr, proc.returncode
    except subprocess.TimeoutExpired:
        return None, f"timed out after {timeout}s", -2
    except OSError as exc:
        return None, str(exc), -3


def _read_file_lines(path, tail=None):
    """Read lines from a file; optionally return only the last *tail* lines."""
    try:
        with open(path, "r", errors="replace") as fh:
            lines = fh.readlines()
        return lines[-tail:] if tail else lines
    except OSError:
        return []

# ---------------------------------------------------------------------------
# 1. Detect GPU devices per architecture via lspci
# ---------------------------------------------------------------------------
def detect_pci_gpus():
    """
    Parse lspci output and classify GPU devices by vendor.

    Returns:
        gpu_map: dict  { arch: [lspci_line, ...] }   per-vendor PCI lines
        pci_gpu_bdfs: dict  { arch: ["00:03.0", ...] }  PCI BDF list per vendor

    Detection rules (matching the project's bash conventions):
        NVIDIA : lspci | grep -E '(3D|VGA compatible) controller' | grep -i nvidia
        PPU    : lspci | grep -E '(3D|VGA compatible) controller' | grep -iE 'Alibaba|1ded:6001'
        AMD    : lspci | grep 'Processing accelerators: Advanced Micro Devices, Inc.'
        NPU    : lspci | grep 'Device 6000' | grep 'Non-VGA'
        Intel  : lspci | grep -E '(VGA|Display) controller' | grep -i intel
    """
    stdout, _, _ = _run("lspci 2>/dev/null")
    if not stdout:
        return {}, {}

    all_lines = stdout.splitlines()
    gpu_map = {}       # arch -> [lspci_line, ...]
    pci_gpu_bdfs = {}  # arch -> ["00:03.0", ...]

    for arch, primary_re, vendor_re in LSPCI_VENDOR_RULES:
        # First pass: filter by device class
        primary_hits = [l for l in all_lines if re.search(primary_re, l, re.I)]
        # Second pass: filter by vendor (if applicable)
        if vendor_re is not None:
            matched = [l for l in primary_hits if re.search(vendor_re, l, re.I)]
        else:
            matched = primary_hits

        if matched:
            gpu_map[arch] = matched
            # Extract PCI BDF from each matched line (e.g. "00:03.0" or "0000:00:03.0")
            bdfs = []
            for line in matched:
                m = re.match(r"\s*([0-9a-fA-F]{4}:)?([0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9a-fA-F])", line)
                if m:
                    bdf = m.group(1, 2)
                    bdfs.append((bdf[0] or "") + bdf[1])
                else:
                    bdfs.append(line.strip().split()[0])
            pci_gpu_bdfs[arch] = bdfs

    return gpu_map, pci_gpu_bdfs

# ---------------------------------------------------------------------------
# Map PCI BDF to GPU index via SMI (NVIDIA)
# Returns: dict { "0000:00:03.0": 0, "0000:84:00.0": 1, ... }
# ---------------------------------------------------------------------------
def _map_pci_to_gpu_index(smi):
    """Query SMI for PCI bus ID -> GPU index mapping."""
    if smi != "nvidia-smi":
        return {}
    stdout, _, rc = _run(
        f"{smi} --query-gpu=index,pci.bus_id --format=csv,noheader",
        timeout=SMI_TIMEOUT,
    )
    if not stdout or rc != 0:
        return {}
    mapping = {}
    for line in stdout.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) >= 2:
            try:
                idx = int(parts[0])
                raw_bus_id = parts[1]   # e.g. "00000000:00:03.0"
                mapping[raw_bus_id] = idx
                # Also index by short form "00:03.0" and "0000:00:03.0"
                # nvidia-smi uses 8-digit domain: "00000000:BB:DD.F"
                # lspci may produce "BB:DD.F" or "DDDD:BB:DD.F"
                m = re.match(
                    r"(?:[0-9a-fA-F]+:)?([0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9a-fA-F])",
                    raw_bus_id,
                )
                if m:
                    short = m.group(1)                      # "00:03.0"
                    mapping[short] = idx
                    mapping[f"0000:{short}"] = idx          # "0000:00:03.0"
            except (ValueError, TypeError):
                continue
    return mapping


def build_gpu_device_entries(arch, pci_bdfs, smi=None):
    """
    Build gpu_devices list in the format: "PCI_BDF(GPU:index)"

    For NVIDIA, cross-references PCI BDF with nvidia-smi to get GPU index.
    For other architectures, uses PCI BDF alone.

    Returns: list of strings like ["0000:00:03.0(GPU:0)", "0000:84:00.0(GPU:1)"]
    """
    if not pci_bdfs:
        return []

    # Try to get PCI -> GPU index mapping
    pci_to_idx = {}
    if smi:
        pci_to_idx = _map_pci_to_gpu_index(smi)

    entries = []
    for bdf in pci_bdfs:
        # Try exact match first, then short form, then with 0000: prefix
        short = bdf
        if re.match(r"[0-9a-fA-F]{4}:[0-9a-fA-F]{2}:", bdf):
            # Already has domain prefix — also compute short form
            short = re.sub(r"^[0-9a-fA-F]+:", "", bdf)  # "0000:00:03.0" -> "00:03.0"
        for key in (bdf, short, f"0000:{short}"):
            if key in pci_to_idx:
                entries.append(f"{bdf}(GPU:{pci_to_idx[key]})")
                break
        else:
            entries.append(bdf)
    return entries


# ---------------------------------------------------------------------------
# Check: GPU driver loaded
# ---------------------------------------------------------------------------
def check_driver(arch, pci_gpus):
    if not pci_gpus:
        _report("driver", "ERROR", "No GPU detected via lspci")
        return

    stdout, _, _ = _run("lsmod 2>/dev/null")
    loaded_mods = set()
    if stdout:
        for line in stdout.splitlines():
            name = line.split()[0] if line.split() else ""
            loaded_mods.add(name)

    found_mod = None
    for mod in DRIVER_MODULES.get(arch, DRIVER_MODULES[ARCH_UNKNOWN]):
        if mod in loaded_mods:
            found_mod = mod
            break

    if found_mod:
        _report("driver", "OK", f"Driver module '{found_mod}' is loaded (arch={arch})")
    else:
        _report("driver", "ERROR",
                f"GPU detected via lspci but no driver module loaded "
                f"(arch={arch}, expected one of: "
                f"{DRIVER_MODULES.get(arch, DRIVER_MODULES[ARCH_UNKNOWN])})")

# ---------------------------------------------------------------------------
# Check: SMI tool exists
# ---------------------------------------------------------------------------
def find_smi(arch):
    """Return the first available SMI tool path, or None."""
    for tool in SMI_PREFERRED.get(arch, SMI_PREFERRED[ARCH_UNKNOWN]):
        if shutil.which(tool):
            return tool
    return None


def check_smi_exists(arch):
    expected = SMI_PREFERRED.get(arch, SMI_PREFERRED[ARCH_UNKNOWN])
    if not expected:
        # This architecture has no SMI tool support
        _report("smi_exists", "OK",
                f"SMI tool not supported for arch '{arch}' — "
                "SMI-dependent checks will be skipped")
        return None
    smi = find_smi(arch)
    if smi is None:
        _report("smi_exists", "ERROR",
                f"No SMI tool found (tried: {expected}). "
                "Install the vendor management toolkit.")
        return None
    _report("smi_exists", "OK", f"SMI tool found: {smi}")
    return smi

# ---------------------------------------------------------------------------
# Check: SMI responsive (not hung, no error)
# ---------------------------------------------------------------------------
def check_smi_responsive(smi):
    if smi is None:
        return False
    stdout, stderr, rc = _run(smi, timeout=SMI_TIMEOUT)
    if stdout is None:
        _report("smi_responsive", "ERROR",
                f"'{smi}' did not respond within {SMI_TIMEOUT}s — "
                "possible driver hang or deadlock")
        return False
    if rc != 0:
        _report("smi_responsive", "ERROR",
                f"'{smi}' exited with code {rc}: {stderr.strip()[:200]}")
        return False
    _report("smi_responsive", "OK", f"'{smi}' responded normally")
    return True

# ---------------------------------------------------------------------------
# Check: Card drop — lspci count vs smi count
# ---------------------------------------------------------------------------
def _nvidia_smi_count(smi):
    stdout, _, rc = _run(
        f"{smi} --query-gpu=index --format=csv,noheader",
        timeout=SMI_TIMEOUT,
    )
    if stdout is None:
        return None, "smi timed out"
    if rc != 0:
        return None, f"smi error rc={rc}"
    return len([l for l in stdout.strip().splitlines() if l.strip()]), None


def _amd_smi_count(smi):
    # rocm-smi has no `list` sub-command. Use the default concise table; each
    # GPU appears on its own line and is referenced as "GPU<N>" / "GPU[<N>]"
    # in the output (works for both old and new rocm-smi versions).
    stdout, _, rc = _run(f"{smi} 2>/dev/null", timeout=SMI_TIMEOUT)
    if stdout is None:
        return None, "smi timed out"
    count = len(re.findall(r"GPU\s*\[?\d+\]?", stdout, re.I))
    return count or 0, None


def _ppu_smi_count(smi):
    """Count PPU devices via ppu-smi (Alibaba PPU/alixpu).

    Reference command (from pouch_pid.lua):
        ppu-smi --query-ppu index,serial,utilization.ppu --format=csv
    """
    stdout, _, rc = _run(
        f"{smi} --query-ppu=index --format=csv,noheader",
        timeout=SMI_TIMEOUT,
    )
    if stdout is None:
        return None, "smi timed out"
    if rc != 0:
        return None, f"smi error rc={rc}"
    return len([l for l in stdout.strip().splitlines() if l.strip()]), None


def check_card_drop(arch, smi, pci_count, expected_count=0):
    """Detect GPU card drop.

    If expected_count > 0 (rated/expected GPU quantity provided):
      - lspci_count < expected_count → PCIe-level card drop
      - smi_count   < expected_count → driver/SMI-level card drop
      - otherwise OK
      (lspci vs smi comparison is NOT performed in this mode)

    If expected_count <= 0 (default / not provided):
      - Original logic: compare smi_count < pci_count
    """
    if pci_count <= 0 or smi is None:
        return

    smi_count, err = None, None
    if arch == ARCH_NVIDIA:
        smi_count, err = _nvidia_smi_count(smi)
    elif arch == ARCH_AMD:
        smi_count, err = _amd_smi_count(smi)
    elif arch == ARCH_PPU:
        smi_count, err = _ppu_smi_count(smi)
    else:
        return

    if smi_count is None:
        _report("card_drop", "WARN", f"Could not query GPU count from {smi}: {err}")
        return

    if expected_count and expected_count > 0:
        # --- Rated count mode ---
        if pci_count < expected_count:
            _report("card_drop", "ERROR",
                    f"Card drop at PCIe level! "
                    f"Expected {expected_count} GPU(s) but lspci shows only {pci_count}. "
                    f"Missing {expected_count - pci_count} card(s) from PCIe bus.")
        elif smi_count < expected_count:
            _report("card_drop", "ERROR",
                    f"Card drop at SMI level! "
                    f"Expected {expected_count} GPU(s), lspci shows {pci_count}, "
                    f"but {smi} reports only {smi_count}. "
                    f"Missing {expected_count - smi_count} card(s) from driver.")
        else:
            _report("card_drop", "OK",
                    f"No card drop: expected={expected_count}, "
                    f"lspci={pci_count}, {smi}={smi_count}")
    else:
        # --- Original mode: lspci vs smi ---
        if smi_count < pci_count:
            _report("card_drop", "ERROR",
                    f"Card drop detected! "
                    f"lspci shows {pci_count} GPU(s) but {smi} reports {smi_count}. "
                    f"Missing {pci_count - smi_count} card(s).")
        else:
            _report("card_drop", "OK",
                    f"No card drop: lspci={pci_count}, {smi}={smi_count}")

# ---------------------------------------------------------------------------
# Per-Xid classification: (severity, suggestion, message)
# Based on NVIDIA official documentation:
# https://docs.nvidia.com/deploy/xid-errors/archive/index.html
#
# suggestion field (optional): condensed guidance from NVIDIA official docs,
# describing the problem type (HW / Driver / App) and recommended action.
# ---------------------------------------------------------------------------
_XID_TABLE = {
    # Xid: (severity, suggestion, message)

    # --- Memory / Page Fault ---
    13:  ("WARN",
          "HW/Driver/App/Thermal: run cuda-memcheck or nvidia-debugdump; check app for OOB access",
          "Graphics Engine Exception: possible HW/Driver/UserApp/Thermal/FB fault"),
    31:  ("WARN",
          "HW/Driver/App: likely app OOB memory access; run cuda-memcheck to verify",
          "GPU memory page fault: likely operator OOB or driver/HW memory error"),

    # --- GPU Stopped / Compute Hang ---
    43:  ("WARN",
          "Driver/App: GPU channel stalled; reset the affected process or driver",
          "GPU stopped processing: Driver or UserApp error"),

    # --- Preemptive Cleanup (cascaded from DBE/fatal) ---
    45:  ("ERROR",
          "Driver: cascaded cleanup after prior fatal error (typically Xid 48 DBE); reset GPU",
          "Preemptive cleanup due to previous errors (typically follows Double-Bit ECC)"),

    # --- ECC Hardware Errors ---
    48:  ("ERROR",
          "HW: uncorrectable double-bit ECC error; retire affected memory page, consider GPU replacement",
          "Double-Bit ECC Error: uncorrectable HW memory fault"),
    63:  ("WARN",
          "HW/Driver: ECC page retirement event recorded; monitor page retirement count",
          "ECC page retirement or row remapping recording event"),
    64:  ("ERROR",
          "HW/Driver: ECC page retirement recording FAILED; manual intervention required",
          "ECC page retirement or row remapper recording FAILURE"),
    92:  ("ERROR",
          "HW: high single-bit ECC rate indicates memory degradation; take GPU offline for page retirement",
          "High single-bit ECC error rate: HW degradation, GPU offline recommended"),
    94:  ("ERROR",
          "HW: contained ECC error, current workload affected; reset GPU to recover",
          "Contained ECC error: HW fault, workload affected but system stable"),
    95:  ("ERROR",
          "HW: uncontained ECC error, system memory integrity compromised; reboot and inspect hardware",
          "Uncontained ECC error: HW fault, system integrity compromised"),

    # --- Micro-controller ---
    61:  ("WARN",
          "Driver: informational micro-controller breakpoint; no action required unless recurring",
          "Internal micro-controller breakpoint/warning (informational)"),
    62:  ("ERROR",
          "HW/Driver/Thermal: micro-controller halted; check thermals, reset GPU or update driver",
          "Internal micro-controller halt: HW/Driver/Thermal fault"),

    # --- Engine / Decoder Exceptions (Driver) ---
    65:  ("WARN",
          "Driver: video processor exception; check driver version and workload",
          "Video processor exception (Driver error)"),
    68:  ("WARN",
          "Driver: NVDEC exception; check decoder workload or update driver",
          "NVDEC0 exception (Driver error)"),
    69:  ("WARN",
          "Driver: graphics engine class error; check driver version",
          "Graphics engine class error (Driver error)"),
    70:  ("WARN",
          "Driver: copy engine exception; check driver version and memory bandwidth",
          "CE (Copy Engine) exception (Driver error)"),

    # --- NVLink ---
    74:  ("ERROR",
          "HW/Driver: NVLink fault; check cable/connector and run nvidia-smi nvlink -s",
          "NVLink error: HW or Driver fault on NVLink interconnect"),

    # --- Card Lost (PCIe) ---
    79:  ("ERROR",
          "HW: GPU lost from PCIe bus; reseat card, check PCIe slot and power supply",
          "GPU has fallen off the PCIe bus: HW fault, card no longer accessible"),

    # --- Data Corruption / Engine Reset ---
    80:  ("ERROR",
          "HW/Driver: data corruption detected; reset GPU via nvidia-smi --gpu-reset",
          "Corrupted data detected: engine reset required (Driver/HW error)"),

    # --- GSP (GPU System Processor) ---
    119: ("ERROR",
          "Driver: GSP RPC timed out; update driver or reset GPU",
          "GSP RPC Timeout: Driver-GSP communication failure"),
    120: ("ERROR",
          "Driver: GSP firmware error; update driver or reset GPU",
          "GSP internal error: Driver/firmware fault on GPU System Processor"),
}

# Fallback for Xid numbers not in _XID_TABLE
_XID_UNKNOWN_ENTRY = ("WARN",
                      "Unknown Xid: collect nvidia-bug-report.sh and contact NVIDIA support",
                      "Unknown Xid error")

# ---------------------------------------------------------------------------
# Compression knobs for repeated Xid records
# ---------------------------------------------------------------------------
_XID_RECENT_HOURS    = 1     # only keep records within the last N hours when timestamped
_XID_MAX_LINES       = 10    # at most N latest lines per Xid number
# ISO timestamp prefix matcher; supports `dmesg --time-format iso` output:
#   2026-04-01T10:00:00,123+08:00 ...   (note the comma fraction separator)
#   2026-04-01T10:00:00+08:00 ...
_ISO_TS_RE = re.compile(
    r"(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})"
    r"(?:[.,]\d+)?"
    r"(Z|[+-]\d{2}:?\d{2})?"
)


def _parse_iso_ts(line):
    """Best-effort parse of a leading ISO timestamp; returns aware datetime or None."""
    m = _ISO_TS_RE.search(line)
    if not m:
        return None
    base, tz = m.group(1), m.group(2)
    if tz:
        if tz == "Z":
            tz = "+0000"
        else:
            # Remove colon: "+08:00" -> "+0800"; keep "+0800" as-is
            tz = tz.replace(":", "")
        try:
            return datetime.datetime.strptime(base + tz, "%Y-%m-%dT%H:%M:%S%z")
        except ValueError:
            return None
    else:
        try:
            dt = datetime.datetime.strptime(base, "%Y-%m-%dT%H:%M:%S")
            return dt.astimezone()  # assume local timezone
        except ValueError:
            return None


def _compress_xid_lines(lines, hours=_XID_RECENT_HOURS, limit=_XID_MAX_LINES):
    """Keep at most `limit` recent lines per Xid.

    For lines carrying a parseable timestamp, only those within the last
    `hours` hours are retained. Lines without a parseable timestamp are
    used as a fallback only when no timestamped lines remain.
    """
    cutoff = datetime.datetime.now().astimezone() - datetime.timedelta(hours=hours)
    timed, untimed = [], []
    for ln in lines:
        ts = _parse_iso_ts(ln)
        if ts is None:
            untimed.append(ln)
        elif ts >= cutoff:
            timed.append(ln)
    pool = timed if timed else untimed
    return pool[-limit:]


# ---------------------------------------------------------------------------
# Check: Xid errors (NVIDIA)
# ---------------------------------------------------------------------------
def check_xid_errors():
    xid_re = re.compile(r"Xid\s*\(PCI:[^)]+\):\s*(\d+)", re.I)
    # Collect per-xid entries: { xid_num: [line, ...] }
    xid_map = {}

    # dmesg
    stdout, _, _ = _run(
        "dmesg --time-format iso 2>/dev/null || dmesg 2>/dev/null",
        timeout=15,
    )
    if stdout:
        for line in stdout.splitlines()[-DMESG_LINES:]:
            m = xid_re.search(line)
            if m:
                xid_num = int(m.group(1))
                xid_map.setdefault(xid_num, []).append(f"[dmesg] {line.strip()}")

    # Persistent log files
    for lf in LOG_FILES:
        if not os.path.isfile(lf):
            continue
        out, _, _ = _run(f"grep -aE 'Xid' {lf} 2>/dev/null | tail -200", timeout=10)
        if out:
            for line in out.splitlines():
                if not line.strip():
                    continue
                m = xid_re.search(line)
                xid_num = int(m.group(1)) if m else -1
                xid_map.setdefault(xid_num, []).append(f"[{lf}] {line.strip()}")

    if not xid_map:
        _report("xid_errors", "OK", "No Xid errors found in dmesg / logs")
        return

    xid_list = sorted(xid_map.keys())

    for xid_num in xid_list:
        # Per-Xid compression: only recent (1h) and at most 10 lines
        lines = _compress_xid_lines(xid_map[xid_num])
        if not lines:
            continue
        sev, suggestion, msg = _XID_TABLE.get(xid_num, _XID_UNKNOWN_ENTRY)
        sample = "\n  ".join(lines)
        _report(
            f"Xid{xid_num}", sev,
            f"{msg} ({len(lines)} recent occurrence(s) within {_XID_RECENT_HOURS}h):\n  {sample}",
            suggestion=suggestion,
        )

# ---------------------------------------------------------------------------
# Check: NVIDIA SMI detailed metrics
# Returns: set of PIDs currently using GPU
# ---------------------------------------------------------------------------
def check_nvidia_metrics(smi, arch):
    if smi != "nvidia-smi":
        return set()

    fields = ",".join([
        "index", "name",
        "ecc.errors.corrected.volatile.total",
        "ecc.errors.uncorrected.volatile.total",
        "temperature.gpu",
        "power.draw", "power.limit",
        "utilization.gpu",
    ])
    stdout, _, rc = _run(
        f"{smi} --query-gpu={fields} --format=csv,noheader,nounits",
        timeout=SMI_TIMEOUT,
    )
    if not stdout or rc != 0:
        return set()

    for line in stdout.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 8:
            continue

        idx, name            = parts[0], parts[1]
        ecc_corr, ecc_uncorr = parts[2], parts[3]
        temp                 = parts[4]
        power_draw, power_lim = parts[5], parts[6]
        gpu_util             = parts[7]
        tag                  = f"GPU[{idx}]({name})"

        # ECC errors
        try:
            uncorr, corr = int(ecc_uncorr), int(ecc_corr)
            if uncorr > 0:
                _report("ecc", "ERROR",
                        f"{tag} Uncorrected ECC errors: {uncorr} — "
                        "hardware memory fault, GPU may need replacement")
            elif corr > 0:
                _report("ecc", "WARN",
                        f"{tag} Corrected ECC errors: {corr} — "
                        "memory degradation, monitor closely")
            else:
                _report("ecc", "OK", f"{tag} No ECC errors")
        except (ValueError, TypeError):
            pass

        # Temperature
        try:
            t = int(temp)
            if t >= 90:
                _report("temperature", "ERROR",
                        f"{tag} Critical temperature: {t}°C (≥90°C) — "
                        "risk of thermal shutdown")
            elif t >= 80:
                _report("temperature", "WARN",
                        f"{tag} High temperature: {t}°C (≥80°C) — "
                        "check cooling system")
            else:
                _report("temperature", "OK", f"{tag} Temperature: {t}°C")
        except (ValueError, TypeError):
            pass

        # Power
        try:
            pd, pl = float(power_draw), float(power_lim)
            if pl > 0:
                ratio = pd / pl * 100
                if ratio >= 95:
                    _report("power", "WARN",
                            f"{tag} Power near limit: {pd:.1f}W / {pl:.1f}W "
                            f"({ratio:.0f}%) — possible throttling")
                else:
                    _report("power", "OK",
                            f"{tag} Power: {pd:.1f}W / {pl:.1f}W ({ratio:.0f}%)")
        except (ValueError, TypeError):
            pass

    # Collect GPU-holding PIDs
    gpu_pids = set()
    stdout2, _, rc2 = _run(
        f"{smi} --query-compute-apps=pid,used_memory --format=csv,noheader,nounits",
        timeout=SMI_TIMEOUT,
    )
    if stdout2 and rc2 == 0:
        for line in stdout2.strip().splitlines():
            p = line.split(",")[0].strip()
            if p.isdigit():
                gpu_pids.add(int(p))

    return gpu_pids

# ---------------------------------------------------------------------------
# Check: Z/D state processes holding GPU
# ---------------------------------------------------------------------------
def check_zombie_disk_on_gpu(gpu_pids):
    issues = []
    proc_root = "/proc"

    if not gpu_pids:
        return

    for pid in sorted(gpu_pids):
        status_path = f"{proc_root}/{pid}/status"
        if not os.path.exists(status_path):
            # Process already gone — not an issue
            continue
        try:
            content = open(status_path, errors="replace").read()
        except OSError:
            continue

        state_m = re.search(r"^State:\s+(\S)", content, re.M)
        name_m  = re.search(r"^Name:\s+(\S+)", content, re.M)
        if not state_m:
            continue

        state = state_m.group(1)
        name  = name_m.group(1) if name_m else "unknown"

        if state == "Z":
            issues.append(
                f"PID {pid} ({name}) is Zombie (Z) but holding GPU — "
                "orphaned process not cleaned up"
            )
        elif state == "D":
            issues.append(
                f"PID {pid} ({name}) is in uninterruptible sleep (D) while holding GPU — "
                "possible driver deadlock or IO hang"
            )

    if issues:
        _report("zombie_on_gpu", "ERROR",
                f"Z/D state processes holding GPU:\n  " + "\n  ".join(issues))
    else:
        _report("zombie_on_gpu", "OK",
                "No Z/D state processes found among GPU-holding processes")

# ---------------------------------------------------------------------------
# Check: dmesg GPU driver error patterns
# ---------------------------------------------------------------------------
def check_dmesg_gpu():
    stdout, _, _ = _run(
        "dmesg --time-format iso 2>/dev/null || dmesg 2>/dev/null",
        timeout=15,
    )
    if not stdout:
        _report("dmesg_gpu", "WARN", "Could not read dmesg output")
        return

    lines = stdout.splitlines()[-DMESG_LINES:]
    # findings: {(severity, desc): [line, ...]}
    findings = {}
    for line in lines:
        for pattern, severity, desc in DMESG_PATTERNS:
            if re.search(pattern, line, re.I):
                key = (severity, desc)
                findings.setdefault(key, []).append(line.strip())

    if not findings:
        _report("dmesg_gpu", "OK", "No GPU driver error patterns found in dmesg")
        return

    for (severity, desc), matched in sorted(findings.items(),
                                            key=lambda x: (x[0][0], x[0][1])):
        sample = matched[-1][:200]
        _report("dmesg_gpu", severity,
                f"{desc} ({len(matched)} occurrence(s)), latest: {sample}")

# ---------------------------------------------------------------------------
# Check: GPU utilization with no registered compute process (process leak)
# ---------------------------------------------------------------------------
def check_gpu_util_no_process(smi, gpu_pids):
    if smi != "nvidia-smi":
        return

    stdout, _, rc = _run(
        f"{smi} --query-gpu=index,utilization.gpu --format=csv,noheader,nounits",
        timeout=SMI_TIMEOUT,
    )
    if not stdout or rc != 0:
        return

    for line in stdout.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 2:
            continue
        try:
            idx, util = parts[0], int(parts[1])
        except (ValueError, TypeError):
            continue

        if util > 5 and not gpu_pids:
            _report("util_leak", "WARN",
                    f"GPU[{idx}] utilization={util}% but no compute process "
                    "registered by nvidia-smi — possible zombie/leaked context")

# ---------------------------------------------------------------------------
# Check: NVLink status (NVIDIA)
# ---------------------------------------------------------------------------
def check_nvlink(smi):
    if smi != "nvidia-smi":
        return

    stdout, _, rc = _run(f"{smi} nvlink --status 2>/dev/null", timeout=SMI_TIMEOUT)
    if not stdout or rc != 0:
        return  # NVLink not present or not supported — not an error

    if re.search(r"inactive|error|fault|fail", stdout, re.I):
        _report("nvlink", "WARN",
                "NVLink anomaly detected — run 'nvidia-smi nvlink --status' for details")
    else:
        _report("nvlink", "OK", "NVLink status normal")

# ---------------------------------------------------------------------------
# JSON output builder
# ---------------------------------------------------------------------------
def build_json_output(gpu_map, gpu_devices_list):
    """
    Build a structured JSON report from accumulated _results.

    The tool intentionally does NOT make any health/anomaly judgement here;
    it just reports every observation. Consumers can decide what to do with
    the per-item ``severity`` field (OK / WARN / ERROR).

    Top-level schema:
    {
      "分析详情": ["<log line>", ...],
      "结果": {
        "timestamp":   "<ISO8601>",
        "hostname":    "<str>",
        "gpu_arch":    ["nvidia", ...],
        "gpu_count":   <int>,
        "gpu_devices": ["00:03.0(GPU:0)", ...],
        "details": [
          {
            "check":      "<check_id>",
            "severity":   "OK|WARN|ERROR",
            "message":    "<str>",
            "suggestion": "<str optional>"
          },
          ...
        ]
      }
    }
    """
    total_count = sum(len(v) for v in gpu_map.values())
    arch_list = sorted(gpu_map.keys())

    result = {
        "timestamp":   datetime.datetime.now().astimezone().isoformat(),
        "hostname":    os.uname().nodename,
        "gpu_arch":    arch_list,
        "gpu_count":   total_count,
        "gpu_devices": gpu_devices_list,
        "details":     list(_results),   # report every observation as-is
    }

    return {
        "分析详情": list(_logs),   # full diagnostic process log
        "结果":    result,
    }


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def parse_args():
    p = argparse.ArgumentParser(
        description="GPU Diagnostic Tool \u2014 detects common GPU anomalies"
    )
    p.add_argument("-f", "--output-file", metavar="FILE",
                   help="Write JSON diagnostic report to FILE (default: stdout)")
    p.add_argument("-n", "--expected-gpu-count", type=int, default=0,
                   metavar="N",
                   help="Rated/expected total GPU count. "
                        "If > 0, card-drop check compares against this value "
                        "instead of lspci vs smi. (default: 0 = disabled)")
    return p.parse_args()


def main():
    args = parse_args()

    if os.geteuid() != 0:
        print("[WARN] Not running as root; some checks (dmesg, /proc) may be limited",
              file=sys.stderr)

    # ── Discover GPUs per architecture ─────────────────────────────────────
    gpu_map, pci_gpu_bdfs = detect_pci_gpus()
    total_count = sum(len(v) for v in gpu_map.values())

    if total_count == 0:
        report = {
            "分析详情": ["未检测到GPU设备"],
            "结果": {
                "timestamp":   datetime.datetime.now().astimezone().isoformat(),
                "hostname":    os.uname().nodename,
                "gpu_arch":    [],
                "gpu_count":   0,
                "gpu_devices": [],
                "details":     [],
            },
        }
        _output_json(report, args.output_file)
        sys.exit(0)

    arch_list = sorted(gpu_map.keys())

    _log(f"[INFO] {total_count} GPU(s) found via lspci, architecture(s): {', '.join(arch_list)}")
    for arch, lines in gpu_map.items():
        for g in lines:
            _log(f"       [{arch}] {g.strip()}")
    _log("")

    # ── Collect gpu_devices entries (PCI BDF + GPU index) ──────────────────
    all_gpu_device_entries = []
    arch_smi = {}      # arch -> smi_tool or None
    arch_smi_ok = {}   # arch -> bool
    arch_gpu_pids = {} # arch -> set of PIDs

    # ── Run checks per architecture ─────────────────────────────────────────
    for arch in arch_list:
        pci_gpus = gpu_map[arch]
        pci_count = len(pci_gpus)
        pci_bdfs = pci_gpu_bdfs.get(arch, [])

        _log(f"── {arch.upper()} checks ({pci_count} device(s)) ──")

        check_driver(arch, pci_gpus)

        # SMI exists
        smi = check_smi_exists(arch)
        arch_smi[arch] = smi

        # SMI responsive
        smi_ok = check_smi_responsive(smi)
        arch_smi_ok[arch] = smi_ok

        # Card drop (per-arch lspci count vs smi count / expected count)
        if smi is not None:
            check_card_drop(arch, smi if smi_ok else None, pci_count,
                            expected_count=args.expected_gpu_count)
        else:
            _report("card_drop", "OK",
                    f"Card drop check not supported for arch '{arch}' (no SMI tool)")

        # Build gpu_devices entries for this arch
        device_entries = build_gpu_device_entries(arch, pci_bdfs, smi if smi_ok else None)
        all_gpu_device_entries.extend(device_entries)

        # NVIDIA-specific metrics
        gpu_pids = set()
        if smi_ok:
            gpu_pids = check_nvidia_metrics(smi, arch) or set()
        arch_gpu_pids[arch] = gpu_pids

        # Z/D state processes
        check_zombie_disk_on_gpu(gpu_pids)

        # Util leak / NVLink (SMI-dependent)
        if smi_ok:
            check_gpu_util_no_process(smi, gpu_pids)
            check_nvlink(smi)

        _log("")

    # ── Global checks (not arch-specific) ───────────────────────────────────
    check_xid_errors()
    check_dmesg_gpu()

    # ── Build and output JSON ───────────────────────────────────────────────
    report = build_json_output(gpu_map, all_gpu_device_entries)
    _output_json(report, args.output_file)

    errors = [r for r in _results if r["severity"] == "ERROR"]
    warns  = [r for r in _results if r["severity"] == "WARN"]
    if errors:
        sys.exit(2)
    if warns:
        sys.exit(1)
    sys.exit(0)


def _output_json(report, output_file=None):
    """Write the JSON report to file or stdout."""
    json_str = json.dumps(report, ensure_ascii=False, indent=2)
    if output_file:
        try:
            with open(output_file, "w", encoding="utf-8") as fh:
                fh.write(json_str)
                fh.write("\n")
        except OSError as exc:
            print(f"[ERROR] Failed to write JSON report to '{output_file}': {exc}",
                  file=sys.stderr)
            sys.exit(2)
    else:
        print(json_str)


if __name__ == "__main__":
    main()
