#!/usr/bin/env python3
"""Opt-in Compute connection diagnostics. Never accepts or executes remote jobs.

Requires Python 3.10+ and, for connection reports, Linux with NVIDIA drivers.
Uses only the Python standard library. Reads inventory with fixed read-only
nvidia-smi queries; does not install software, change clocks/power, access
wallets, or mine. `diagnose` also prints a local rig report: GPU groups, PCIe
links and rule-of-thumb use-case verdicts. It is never sent to Compute.
"""
from __future__ import annotations

import argparse
import csv
import io
import json
import math
import os
from pathlib import Path
import platform
import re
import shutil
import stat
import subprocess
import sys
import time
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import urlsplit, urlunsplit
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener

VERSION = "0.1.0"
TOKEN_ENV = "COMPUTE_NODE_TOKEN"
TOKEN_PATTERN = re.compile(r"cpn_[A-Za-z0-9_-]{43}\Z")
CHALLENGE_PATTERN = re.compile(r"[A-Za-z0-9_-]{43}\Z")
MAX_RESPONSE = 65_536
HEARTBEAT_SECONDS = 30


class WorkerError(Exception):
    """An actionable error whose message never contains a credential."""


class NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        raise WorkerError("The service redirected this request. Use its final HTTPS endpoint; credentials were not forwarded.")


def endpoint_url(value: str, allow_local_preview: bool = False) -> str:
    try:
        parts = urlsplit(value)
        _ = parts.port
    except ValueError as exc:
        raise WorkerError("Enter a valid worker endpoint URL.") from exc
    if parts.username or parts.password or parts.query or parts.fragment:
        raise WorkerError("The endpoint must not contain credentials, a query, or a fragment.")
    if not parts.hostname or parts.path.rstrip("/") != "/api/providers/worker":
        raise WorkerError("Use the worker endpoint shown when registering the node, ending in /api/providers/worker.")
    loopback = parts.hostname in ("127.0.0.1", "localhost", "::1")
    if parts.scheme != "https" and not (parts.scheme == "http" and loopback and allow_local_preview):
        raise WorkerError("HTTPS is required. Loopback HTTP is allowed only with --allow-local-preview and a preview server.")
    return urlunsplit((parts.scheme, parts.netloc, "/api/providers/worker", "", ""))


def request_json(url: str, token: str | None = None, body: dict[str, Any] | None = None) -> dict[str, Any]:
    headers = {"Accept": "application/json", "User-Agent": f"compute-contributor/{VERSION}"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    encoded = None
    if body is not None:
        encoded = json.dumps(body, separators=(",", ":")).encode("utf-8")
        headers["Content-Type"] = "application/json"
    # Ignore environment proxy configuration; never forward a bearer through an
    # unexpected proxy. Default HTTPS verification remains enabled.
    opener = build_opener(ProxyHandler({}), NoRedirect())
    request = Request(url, data=encoded, headers=headers, method="POST" if body is not None else "GET")
    try:
        with opener.open(request, timeout=15) as response:
            raw = response.read(MAX_RESPONSE + 1)
    except HTTPError as exc:
        if exc.code == 401:
            raise WorkerError("The node credential is invalid, expired, or revoked. Create a new connection credential in your workspace.") from exc
        if exc.code == 409:
            raise WorkerError("The connection challenge expired or was replaced. Retry; ensure only one worker uses this credential.") from exc
        if exc.code == 429:
            raise WorkerError("Too many connection attempts. Wait at least one minute before retrying.") from exc
        raise WorkerError(f"The service returned HTTP {exc.code}. Check its availability and endpoint.") from exc
    except (URLError, TimeoutError, OSError) as exc:
        raise WorkerError("Cannot reach the service with a verified connection. Check its endpoint, certificate, and network access.") from exc
    if len(raw) > MAX_RESPONSE:
        raise WorkerError("The service response exceeded the supported size.")
    try:
        result = json.loads(raw)
    except (ValueError, UnicodeError) as exc:
        raise WorkerError("The service did not return a valid JSON response.") from exc
    if not isinstance(result, dict):
        raise WorkerError("The service returned an unexpected response.")
    return result


def verify_preview(endpoint: str) -> None:
    if urlsplit(endpoint).scheme != "http":
        return
    parts = urlsplit(endpoint)
    status = request_json(urlunsplit((parts.scheme, parts.netloc, "/api/network", "", "")))
    overview = status.get("network", status)
    if not isinstance(overview, dict) or overview.get("mode") != "preview":
        raise WorkerError("Loopback HTTP is permitted only for a server in explicit preview mode. No credential was sent.")


def read_token(token_file: str | None) -> str:
    if token_file:
        path = Path(token_file).expanduser()
        try:
            metadata = path.lstat()
            if not stat.S_ISREG(metadata.st_mode) or metadata.st_mode & 0o077:
                raise WorkerError("The credential file must be a regular private file with mode 600. Symlinks are not accepted.")
            if hasattr(os, "getuid") and metadata.st_uid != os.getuid():
                raise WorkerError("The credential file must be owned by the current user.")
            if metadata.st_size > 256:
                raise WorkerError("The credential file is not a valid node credential.")
            token = path.read_text(encoding="utf-8").strip()
        except (OSError, UnicodeError) as exc:
            raise WorkerError("The private credential file could not be read.") from exc
    else:
        token = os.environ.get(TOKEN_ENV, "").strip()
    if not TOKEN_PATTERN.fullmatch(token):
        raise WorkerError(f"Provide the connection credential using --token-file or the {TOKEN_ENV} environment variable. Never pass a wallet private key.")
    return token


def detect_hardware() -> dict[str, Any]:
    blockers: list[str] = []
    system = platform.system().lower()
    detected: list[dict[str, Any]] = []
    if system != "linux":
        blockers.append("The contributor pilot supports Linux with NVIDIA GPUs. This machine is not Linux.")
    executable = shutil.which("nvidia-smi")
    if not executable:
        blockers.append("nvidia-smi was not found. Install and validate NVIDIA drivers yourself before connecting.")
    elif system == "linux":
        try:
            # Fixed arguments only. No shell is used and server data never enters
            # this command. GPU names and physical VRAM are reported, not attested.
            completed = subprocess.run(
                [executable, "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"],
                check=True, capture_output=True, text=True, timeout=10,
            )
            if len(completed.stdout) > 16_384:
                raise WorkerError("NVIDIA inventory output was unexpectedly large.")
            for row in csv.reader(io.StringIO(completed.stdout)):
                if len(row) != 2:
                    raise WorkerError("NVIDIA inventory could not be read. Ensure driver diagnostics work.")
                name = row[0].strip()
                if not name.lower().startswith("nvidia "):
                    name = f"NVIDIA {name}"
                if len(name) > 100 or any(ord(character) < 32 or ord(character) == 127 for character in name):
                    raise WorkerError("NVIDIA inventory returned an unsupported GPU name.")
                memory_mib = float(row[1].strip())
                if not math.isfinite(memory_mib) or memory_mib <= 0:
                    raise WorkerError("NVIDIA inventory returned invalid memory capacity.")
                # Driver-reported capacity is often slightly below the marketed
                # GiB value (e.g. 24564MiB on a 24GiB card); preserve raw MiB locally.
                detected.append({"gpuModel": name, "vramGb": int(round(memory_mib / 1024)), "reportedVramMiB": memory_mib})
        except (subprocess.SubprocessError, OSError, ValueError) as exc:
            blockers.append("nvidia-smi could not provide supported inventory. Run driver diagnostics before connecting.")
        except WorkerError as exc:
            blockers.append(str(exc))
    if not detected and system == "linux" and executable and not blockers:
        blockers.append("No NVIDIA GPUs were detected.")
    if detected:
        if not 1 <= len(detected) <= 8:
            blockers.append("Register between 1 and 8 GPUs per node.")
        if any(not 4 <= device["vramGb"] <= 256 for device in detected):
            blockers.append("Each GPU must report between 4 and 256GiB for this onboarding pilot.")
        if len({(device["gpuModel"], device["vramGb"]) for device in detected}) != 1:
            blockers.append("This worker currently supports matching GPU models and memory capacities within one node.")
    inventory = None
    if detected and not blockers:
        inventory = {"os": "linux", "gpuModel": detected[0]["gpuModel"], "vramGb": detected[0]["vramGb"], "gpuCount": len(detected)}
    return {
        "workerVersion": VERSION, "platform": system, "readyToConnect": not blockers,
        "inventory": inventory, "detectedGpus": detected, "blockers": blockers,
        "hardwareVerified": False, "dispatchEnabled": False,
        "notice": "Diagnostics report hardware only. No jobs, power changes, mining, or earnings are enabled.",
    }


def require_diagnostic_inventory(diagnostic: dict[str, Any]) -> dict[str, Any]:
    inventory = diagnostic.get("inventory")
    if not diagnostic.get("readyToConnect") or not isinstance(inventory, dict):
        raise WorkerError("This machine is not ready to connect. Run the diagnose command and resolve its blockers.")
    return inventory


def collect_gpu_metrics() -> list[dict[str, Any]] | None:
    """Read current optional NVIDIA counters; unavailable readings stay unknown.

    This fixed query never changes clocks/power or runs a GPU workload. Inventory
    detection remains independent so unsupported counters cannot invent devices.
    """
    if platform.system().lower() != "linux":
        return None
    executable = shutil.which("nvidia-smi")
    if not executable:
        return None
    try:
        completed = subprocess.run(
            [executable, "--query-gpu=index,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw", "--format=csv,noheader,nounits"],
            check=True, capture_output=True, text=True, timeout=5,
        )
        if len(completed.stdout) > 8192:
            return None
        devices: list[dict[str, Any]] = []

        def value(raw: str, low: float, high: float, integer: bool = False):
            text = raw.strip()
            if text.lower().strip("[]") in ("n/a", "not supported", "not available"):
                return None
            number = float(text)
            if not math.isfinite(number) or not low <= number <= high or (integer and not number.is_integer()):
                raise ValueError("Unsupported NVIDIA metric")
            return int(number) if integer else number

        for row in csv.reader(io.StringIO(completed.stdout)):
            if len(row) != 6 or len(devices) >= 8:
                return None
            index = value(row[0], 0, 7, True)
            if index is None or any(device["gpuIndex"] == index for device in devices):
                return None
            device = {
                "gpuIndex": index,
                "utilizationPercent": value(row[1], 0, 100),
                "memoryUsedMiB": value(row[2], 0, 262144, True),
                "memoryTotalMiB": value(row[3], 1, 262144, True),
                "temperatureC": value(row[4], -50, 150),
                "powerDrawWatts": value(row[5], 0, 2000),
            }
            if device["memoryUsedMiB"] is not None and device["memoryTotalMiB"] is not None and device["memoryUsedMiB"] > device["memoryTotalMiB"]:
                return None
            devices.append(device)
        if not devices or sorted(device["gpuIndex"] for device in devices) != list(range(len(devices))):
            return None
        return sorted(devices, key=lambda device: device["gpuIndex"])
    except (subprocess.SubprocessError, OSError, ValueError, csv.Error):
        return None


RIG_DETAIL_FIELDS = (
    "index", "name", "memory.total", "driver_version", "compute_cap",
    "pcie.link.gen.current", "pcie.link.gen.max", "pcie.link.width.current", "pcie.link.width.max", "power.limit",
)
MAX_RIG_GPUS = 64
MAX_GPUS_PER_MACHINE = 8
NARROW_PCIE_WIDTH = 8  # x1 and x4 risers are "narrow"; x8 and x16 slots are not.

# Use-case guidance for the local rig report. These are conservative rules of
# thumb, not qualification: an operator still reviews and approves each machine
# for a specific workload, and nothing here is sent to Compute. Memory is per
# GPU because VRAM does not pool across cards, least of all over x1 risers.
#   good_vram_gb     at or above: enough memory per GPU for typical work
#   limited_vram_gb  at or above (and below good): usable with clear limits
#   min_width_good   narrowest PCIe link needed for "good" (None: links do not matter)
#   min_cc_good      lowest CUDA compute capability for "good" (None: not considered)
USE_CASE_RULES: tuple[dict[str, Any], ...] = (
    {"id": "ai_inference", "label": "AI inference", "good_vram_gb": 24, "limited_vram_gb": 12, "min_width_good": None, "min_cc_good": None,
     "good": "Enough memory per GPU to serve most current open models.",
     "limited": "Fits smaller or quantized models; most current large models need 24 GB or more per GPU.",
     "not_suitable": "Under 12 GB per GPU is too little for most current models."},
    {"id": "ai_training", "label": "AI training", "good_vram_gb": 24, "limited_vram_gb": 12, "min_width_good": NARROW_PCIE_WIDTH, "min_cc_good": None,
     "good": "Enough memory per GPU and a wide enough PCIe link for fine-tuning and training jobs.",
     "limited": "Small models and fine-tuning only; training wants 24 GB or more per GPU on x8 or wider links.",
     "narrow": "Enough memory, but narrow PCIe links slow data loading and multi-GPU training; x8 or wider is recommended.",
     "unknown_width": "Enough memory, but the PCIe link width is unknown; training wants x8 or wider.",
     "not_suitable": "Under 12 GB per GPU is too little to train most current models."},
    {"id": "rendering_3d", "label": "3D rendering", "good_vram_gb": 8, "limited_vram_gb": 4, "min_width_good": None, "min_cc_good": None,
     "good": "Enough memory for most scenes; risers are fine because a scene loads once per job.",
     "limited": "Simple scenes only; large scenes and textures need 8 GB or more per GPU.",
     "not_suitable": "Under 4 GB per GPU is too little for current renderers."},
    {"id": "scientific", "label": "Scientific and engineering compute", "good_vram_gb": 8, "limited_vram_gb": 0, "min_width_good": None, "min_cc_good": (7, 0),
     "good": "A modern CUDA architecture with enough memory for most GPU-accelerated solvers.",
     "limited": "Under 8 GB per GPU limits problem size for most GPU-accelerated solvers.",
     "old_cc": "Compute capability below 7.0: an older architecture that newer CUDA software is dropping.",
     "unknown_cc": "The driver did not report compute capability; most current CUDA software needs 7.0 or newer.",
     "not_suitable": "Not suitable."},
)


def _rig_value(raw: str, kind: str) -> Any:
    """Parse one optional nvidia-smi field; unavailable or implausible values become None."""
    text = raw.strip()
    if text.lower().strip("[]") in ("n/a", "not supported", "not available", "unknown error", ""):
        return None
    try:
        if kind == "cc":
            match = re.fullmatch(r"(\d{1,2})\.(\d{1,2})", text)
            return f"{int(match.group(1))}.{int(match.group(2))}" if match else None
        if kind == "driver":
            return text if re.fullmatch(r"[0-9]{1,4}(\.[0-9]{1,4}){0,3}", text) else None
        number = float(text)
        if not math.isfinite(number):
            return None
        if kind == "gen":
            return int(number) if number.is_integer() and 1 <= number <= 7 else None
        if kind == "width":
            return int(number) if number.is_integer() and number in (1, 2, 4, 8, 12, 16, 32) else None
        if kind == "watts":
            return round(number, 2) if 1 <= number <= 2000 else None
        if kind == "mib":
            return int(round(number)) if 1 <= number <= 262_144 else None
    except ValueError:
        return None
    return None


def _gpu_name(raw: str) -> str | None:
    name = raw.strip()
    if not name.lower().startswith("nvidia "):
        name = f"NVIDIA {name}"
    if len(name) > 100 or any(ord(character) < 32 or ord(character) == 127 for character in name):
        return None
    return name


def collect_rig_details() -> list[dict[str, Any]] | None:
    """Read per-GPU driver, compute capability, PCIe link and power-limit fields.

    One fixed read-only query, retried once without compute_cap for drivers
    that predate that field. Any single unavailable field is reported as None.
    """
    if platform.system().lower() != "linux":
        return None
    executable = shutil.which("nvidia-smi")
    if not executable:
        return None
    for fields in (RIG_DETAIL_FIELDS, tuple(field for field in RIG_DETAIL_FIELDS if field != "compute_cap")):
        try:
            completed = subprocess.run(
                [executable, f"--query-gpu={','.join(fields)}", "--format=csv,noheader,nounits"],
                check=True, capture_output=True, text=True, timeout=10,
            )
        except (subprocess.SubprocessError, OSError):
            continue
        if len(completed.stdout) > 32_768:
            return None
        devices: list[dict[str, Any]] = []
        try:
            for row in csv.reader(io.StringIO(completed.stdout)):
                if not row:
                    continue
                if len(row) != len(fields) or len(devices) >= MAX_RIG_GPUS:
                    return None
                values = dict(zip(fields, row))
                index = values["index"].strip()
                name = _gpu_name(values["name"])
                if not index.isdigit() or int(index) >= MAX_RIG_GPUS or name is None:
                    return None
                devices.append({
                    "index": int(index), "name": name,
                    "vramMiB": _rig_value(values["memory.total"], "mib"),
                    "driverVersion": _rig_value(values["driver_version"], "driver"),
                    "computeCapability": _rig_value(values["compute_cap"], "cc") if "compute_cap" in values else None,
                    "pcieGenCurrent": _rig_value(values["pcie.link.gen.current"], "gen"),
                    "pcieGenMax": _rig_value(values["pcie.link.gen.max"], "gen"),
                    "pcieWidthCurrent": _rig_value(values["pcie.link.width.current"], "width"),
                    "pcieWidthMax": _rig_value(values["pcie.link.width.max"], "width"),
                    "powerLimitWatts": _rig_value(values["power.limit"], "watts"),
                })
        except csv.Error:
            return None
        if not devices or len({device["index"] for device in devices}) != len(devices):
            return None
        return sorted(devices, key=lambda device: device["index"])
    return None


def _cc_tuple(value: str | None) -> tuple[int, int] | None:
    if not value:
        return None
    major, minor = value.split(".")
    return int(major), int(minor)


def _index_list(indices: list[int]) -> str:
    if len(indices) > 1 and indices == list(range(indices[0], indices[-1] + 1)):
        return f"GPUs {indices[0]}–{indices[-1]}"
    return ("GPU " if len(indices) == 1 else "GPUs ") + ", ".join(str(index) for index in indices)


def _verdict(rule: dict[str, Any], vram_gb: int, narrowest_width: int | None, compute_cap: tuple[int, int] | None) -> dict[str, str]:
    def result(verdict: str, reason_key: str) -> dict[str, str]:
        return {"id": rule["id"], "label": rule["label"], "verdict": verdict, "reason": rule[reason_key]}

    if vram_gb < rule["limited_vram_gb"]:
        return result("not_suitable", "not_suitable")
    if vram_gb < rule["good_vram_gb"]:
        return result("limited", "limited")
    if rule["min_width_good"] is not None:
        if narrowest_width is None:
            return result("limited", "unknown_width")
        if narrowest_width < rule["min_width_good"]:
            return result("limited", "narrow")
    if rule["min_cc_good"] is not None:
        if compute_cap is None:
            return result("limited", "unknown_cc")
        if compute_cap < rule["min_cc_good"]:
            return result("limited", "old_cc")
    return result("good", "good")


def build_rig_report(diagnostic: dict[str, Any], details: list[dict[str, Any]] | None) -> dict[str, Any]:
    """Explain locally what this rig suits and what must change before it can connect.

    Diagnose output only: never part of the connection heartbeat.
    """
    if details:
        gpus = details
    else:
        # Fall back to the inventory query; extra fields are simply unknown.
        gpus = [{
            "index": position, "name": device["gpuModel"], "vramMiB": int(round(device["reportedVramMiB"])),
            "driverVersion": None, "computeCapability": None, "pcieGenCurrent": None, "pcieGenMax": None,
            "pcieWidthCurrent": None, "pcieWidthMax": None, "powerLimitWatts": None,
        } for position, device in enumerate(diagnostic.get("detectedGpus") or [])]
    grouped: dict[tuple[str, int], list[dict[str, Any]]] = {}
    for gpu in gpus:
        vram_gb = int(round(gpu["vramMiB"] / 1024)) if gpu["vramMiB"] else 0
        grouped.setdefault((gpu["name"], vram_gb), []).append(gpu)
    groups: list[dict[str, Any]] = []
    for (name, vram_gb), members in sorted(grouped.items(), key=lambda item: (-len(item[1]), item[0])):
        indices = [gpu["index"] for gpu in members]
        widths = [gpu["pcieWidthCurrent"] for gpu in members]
        narrowest = None if any(width is None for width in widths) else min(widths)
        max_widths = [gpu["pcieWidthMax"] for gpu in members if gpu["pcieWidthMax"] is not None]
        caps = {gpu["computeCapability"] for gpu in members}
        compute_cap = next(iter(caps)) if len(caps) == 1 else None
        pcie_note = None
        if narrowest is not None and narrowest < NARROW_PCIE_WIDTH:
            narrow = [gpu["index"] for gpu in members if gpu["pcieWidthCurrent"] < NARROW_PCIE_WIDTH]
            supported = f" (the cards support x{max(max_widths)})" if max_widths and max(max_widths) > narrowest else ""
            pcie_note = (f"{_index_list(narrow)} run at PCIe x{narrowest}{supported}, typical of mining risers. "
                         "Risers are slow to load models but fine for rendering, where a scene loads once per job.")
        elif narrowest is None and details:
            pcie_note = "The driver did not report PCIe link width for every GPU in this group."
        fits = 4 <= vram_gb <= 256
        groups.append({
            "gpuModel": name, "vramGb": vram_gb, "count": len(members), "indices": indices,
            "computeCapability": compute_cap, "pcieWidthNarrowest": narrowest,
            "machinesNeeded": math.ceil(len(members) / MAX_GPUS_PER_MACHINE) if fits else 0,
            "useCases": [_verdict(rule, vram_gb, narrowest, _cc_tuple(compute_cap)) for rule in USE_CASE_RULES],
            "pcieNote": pcie_note,
        })
    notes: list[str] = []
    if len(groups) > 1:
        notes.append(
            f"This host mixes {len(groups)} kinds of GPU. Today a machine is listed with one GPU model and memory size, "
            "and the worker reports every GPU on the host, so it cannot connect as one machine. "
            "To list these cards, move them so each Linux host holds a single model."
        )
    if len(gpus) > MAX_GPUS_PER_MACHINE:
        notes.append(
            f"This host has {len(gpus)} GPUs. A machine can register 1 to {MAX_GPUS_PER_MACHINE} GPUs, and the worker "
            f"cannot select a subset on one host, so move cards until each host has {MAX_GPUS_PER_MACHINE} or fewer."
        )
    for group in groups:
        if not group["vramGb"]:
            notes.append(f"{group['gpuModel']} did not report its memory; machines need 4 to 256 GB per GPU.")
        elif not 4 <= group["vramGb"] <= 256:
            notes.append(f"{group['gpuModel']} reports {group['vramGb']} GB; machines need 4 to 256 GB per GPU.")
    return {
        "reportVersion": 1, "detailsAvailable": bool(details), "gpus": gpus, "groups": groups, "notes": notes,
        "guidance": "These verdicts are rules of thumb from fixed memory, PCIe and compute-capability thresholds, "
                    "not qualification. An operator still reviews and approves each machine for a specific workload.",
    }


def connect_once(endpoint: str, token: str, inventory: dict[str, Any]) -> dict[str, Any]:
    metrics = collect_gpu_metrics()
    challenge = request_json(endpoint + "/challenge", token, {})
    nonce = challenge.get("challenge")
    if not isinstance(nonce, str) or not CHALLENGE_PATTERN.fullmatch(nonce):
        raise WorkerError("The service returned an invalid connection challenge.")
    if urlsplit(endpoint).scheme == "http" and challenge.get("mode") != "preview":
        raise WorkerError("This service is not in preview mode. Use HTTPS.")
    result = request_json(endpoint + "/heartbeat", token, {**inventory, "challenge": nonce, "agentVersion": VERSION, **({"gpuMetrics": metrics} if metrics is not None else {})})
    # This release has no executor. Fail closed if connected to an unexpected
    # service version rather than interpreting a new response as instructions.
    if result.get("dispatchEnabled") is not False or not isinstance(result.get("node"), dict):
        raise WorkerError("The service returned an unsupported worker protocol. This release accepts connection diagnostics only.")
    return result


def print_result(value: dict[str, Any]) -> None:
    # JSON escapes control characters from labels and remote strings. Credentials
    # are never part of diagnostics; the response contains no server executable.
    print(json.dumps(value, indent=2, sort_keys=True))


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    result.add_argument("--version", action="version", version=VERSION)
    commands = result.add_subparsers(dest="command", required=True)
    commands.add_parser("diagnose", help="Read local hardware diagnostics without contacting Compute.")
    for name, help_text in (("status", "Read the registered node status."), ("connect", "Send connection reports; never execute GPU jobs.")):
        command = commands.add_parser(name, help=help_text)
        command.add_argument("--endpoint", required=True, help="Worker endpoint shown in your workspace.")
        command.add_argument("--token-file", help="Private mode-600 file containing the node credential; otherwise use COMPUTE_NODE_TOKEN.")
        command.add_argument("--allow-local-preview", action="store_true", help="Permit loopback HTTP only after confirming server preview mode.")
        if name == "connect":
            command.add_argument("--once", action="store_true", help="Send one challenge/heartbeat, then exit.")
            command.add_argument("--dry-run", action="store_true", help="Inspect local readiness only; no credential is read and no network request is made.")
    return result


def main(argv: list[str] | None = None) -> int:
    args = parser().parse_args(argv)
    try:
        if args.command == "diagnose":
            diagnostic = detect_hardware()
            # The rig report is local diagnose output only; connect never sends it.
            print_result({**diagnostic, "gpuMetrics": collect_gpu_metrics(), "rig": build_rig_report(diagnostic, collect_rig_details())})
            return 0 if diagnostic["readyToConnect"] else 2
        endpoint = endpoint_url(args.endpoint, args.allow_local_preview)
        inventory = None
        if args.command == "connect":
            diagnostic = detect_hardware()
            if args.dry_run:
                print_result({**diagnostic, "endpoint": endpoint, "dryRun": True, "networkRequests": 0})
                return 0 if diagnostic["readyToConnect"] else 2
            inventory = require_diagnostic_inventory(diagnostic)
        verify_preview(endpoint)
        token = read_token(args.token_file)
        if args.command == "status":
            print_result(request_json(endpoint + "/status", token))
            return 0
        assert inventory is not None
        while True:
            print_result(connect_once(endpoint, token, inventory))
            if args.once:
                return 0
            time.sleep(HEARTBEAT_SECONDS)
            inventory = require_diagnostic_inventory(detect_hardware())
    except KeyboardInterrupt:
        print("Connection reporting stopped. No jobs were running.", file=sys.stderr)
        return 130
    except WorkerError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())
