from __future__ import annotations

import platform
import shutil
import subprocess

import torch


def run_command(cmd: list[str]) -> str:
    if shutil.which(cmd[0]) is None:
        return "not found"
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=False)
        text = (result.stdout or result.stderr).strip()
        return text or f"command exited with code {result.returncode}"
    except Exception as exc:  # pragma: no cover - defensive
        return f"error: {exc}"


def main() -> None:
    print("=== System ===")
    print("python:", platform.python_version())
    print("platform:", platform.platform())

    print("\n=== PyTorch ===")
    print("torch:", torch.__version__)
    print("torch.cuda.is_available():", torch.cuda.is_available())
    print("torch.version.cuda:", torch.version.cuda)

    if torch.cuda.is_available():
        device_count = torch.cuda.device_count()
        print("cuda device count:", device_count)
        for idx in range(device_count):
            name = torch.cuda.get_device_name(idx)
            capability = torch.cuda.get_device_capability(idx)
            print(f"device {idx}: {name} | capability={capability[0]}.{capability[1]}")
    else:
        print("no CUDA device visible to PyTorch")

    print("\n=== Triton ===")
    try:
        import triton  # type: ignore

        print("triton:", triton.__version__)
    except Exception as exc:
        print("triton import failed:", exc)

    print("\n=== Toolkit / Driver Hints ===")
    print("nvcc --version:")
    print(run_command(["nvcc", "--version"]))
    print("\nnvidia-smi:")
    print(run_command(["nvidia-smi"]))


if __name__ == "__main__":
    main()