#!/usr/bin/env python3
# vk_probe.py - enumerate real Vulkan physical devices via the loader (ctypes,
# no dev headers needed). Prints one DEV line per physical device.
import ctypes
import sys

VK_SUCCESS = 0
VK_STRUCTURE_TYPE_APPLICATION_INFO = 0
VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO = 1
VK_API_VERSION_1_2 = (1 << 22) | (2 << 12)

try:
    vulkan = ctypes.CDLL("libvulkan.so.1")
except OSError as e:
    print("LOAD_FAIL %s" % e)
    sys.exit(1)


class VkApplicationInfo(ctypes.Structure):
    _fields_ = [
        ("sType", ctypes.c_int),
        ("pNext", ctypes.c_void_p),
        ("pApplicationName", ctypes.c_char_p),
        ("applicationVersion", ctypes.c_uint32),
        ("pEngineName", ctypes.c_char_p),
        ("engineVersion", ctypes.c_uint32),
        ("apiVersion", ctypes.c_uint32),
    ]


class VkInstanceCreateInfo(ctypes.Structure):
    _fields_ = [
        ("sType", ctypes.c_int),
        ("pNext", ctypes.c_void_p),
        ("flags", ctypes.c_uint32),
        ("pApplicationInfo", ctypes.POINTER(VkApplicationInfo)),
        ("enabledLayerCount", ctypes.c_uint32),
        ("ppEnabledLayerNames", ctypes.c_void_p),
        ("enabledExtensionCount", ctypes.c_uint32),
        ("ppEnabledExtensionNames", ctypes.c_void_p),
    ]


app = VkApplicationInfo(
    VK_STRUCTURE_TYPE_APPLICATION_INFO, None, b"vk_probe", 1, b"vk_probe", 1,
    VK_API_VERSION_1_2)
ci = VkInstanceCreateInfo(
    VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO, None, 0, ctypes.pointer(app), 0,
    None, 0, None)
inst = ctypes.c_void_p()
r = vulkan.vkCreateInstance(ctypes.byref(ci), None, ctypes.byref(inst))
print("vkCreateInstance rc=%d (0x%x)" % (r, r & 0xfffffffff if r < 0 else r))
if r != VK_SUCCESS:
    sys.exit(1)

n = ctypes.c_uint32(0)
r = vulkan.vkEnumeratePhysicalDevices(inst, ctypes.byref(n), None)
print("vkEnumeratePhysicalDevices rc=%d count=%d" % (r, n.value))
if r != VK_SUCCESS or n.value == 0:
    sys.exit(2)

devs = (ctypes.c_void_p * n.value)()
r = vulkan.vkEnumeratePhysicalDevices(inst, ctypes.byref(n), devs)
ok = 0
for i in range(n.value):
    buf = ctypes.create_string_buffer(4096)
    vulkan.vkGetPhysicalDeviceProperties(devs[i], buf)
    raw = buf.raw
    api = int.from_bytes(raw[0:4], "little")
    drv = int.from_bytes(raw[4:8], "little")
    vendor = int.from_bytes(raw[8:12], "little")
    devid = int.from_bytes(raw[12:16], "little")
    dtype = int.from_bytes(raw[16:20], "little")
    name = raw[20:20 + 256].split(b"\x00")[0].decode("utf-8", "replace")
    print("DEV %d: %s vendor=0x%04x device=0x%04x type=%d api=%d.%d.%d drv=%d"
          % (i, name, vendor, devid, dtype, api >> 22,
             (api >> 12) & 0x3ff, api & 0xfff, drv))
    if vendor == 0x10de:
        ok = 1
print("VK_NVIDIA=%d" % ok)
sys.exit(0 if ok else 3)