#!/usr/bin/env python3
# vk_probe2.py - Vulkan probe stage 2: create a logical device + queue on the
# NVIDIA physical device. Prints the queue-family table and the create
# result. ctypes only, no dev headers.
import ctypes
import sys

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

vulkan = ctypes.CDLL("libvulkan.so.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),
    ]


class VkQueueFamilyProperties(ctypes.Structure):
    _fields_ = [
        ("queueFlags", ctypes.c_uint32),
        ("queueCount", ctypes.c_uint32),
        ("timestampValidBits", ctypes.c_uint32),
        ("minImageTransferGranularity", ctypes.c_uint8 * 12),
    ]


class VkDeviceQueueCreateInfo(ctypes.Structure):
    _fields_ = [
        ("sType", ctypes.c_int),
        ("pNext", ctypes.c_void_p),
        ("flags", ctypes.c_uint32),
        ("queueFamilyIndex", ctypes.c_uint32),
        ("queueCount", ctypes.c_uint32),
        ("pQueuePriorities", ctypes.POINTER(ctypes.c_float)),
    ]


class VkDeviceCreateInfo(ctypes.Structure):
    _fields_ = [
        ("sType", ctypes.c_int),
        ("pNext", ctypes.c_void_p),
        ("flags", ctypes.c_uint32),
        ("queueCreateInfoCount", ctypes.c_uint32),
        ("pQueueCreateInfos", ctypes.POINTER(VkDeviceQueueCreateInfo)),
        ("enabledLayerCount", ctypes.c_uint32),
        ("ppEnabledLayerNames", ctypes.c_void_p),
        ("enabledExtensionCount", ctypes.c_uint32),
        ("ppEnabledExtensionNames", ctypes.c_void_p),
        ("pEnabledFeatures", ctypes.c_void_p),
    ]


app = VkApplicationInfo(VK_STRUCTURE_TYPE_APPLICATION_INFO, None, b"vk_probe2",
                        1, b"vk_probe2", 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("INSTANCE rc=%d" % r)
if r != 0:
    sys.exit(1)

n = ctypes.c_uint32(0)
vulkan.vkEnumeratePhysicalDevices(inst, ctypes.byref(n), None)
devs = (ctypes.c_void_p * n.value)()
vulkan.vkEnumeratePhysicalDevices(inst, ctypes.byref(n), devs)

# pick the NVIDIA device by name
best = None
for i in range(n.value):
    buf = ctypes.create_string_buffer(4096)
    vulkan.vkGetPhysicalDeviceProperties(devs[i], buf)
    raw = buf.raw
    name = raw[20:276].split(b"\x00")[0].decode("utf-8", "replace")
    print("PROPS dev=%d head=%s name=%r" % (i, raw[:40].hex(), name))
    if "NVIDIA" in name or "llvmpipe" in name:
        if "NVIDIA" in name:
            best = devs[i]
            print("PICK NVIDIA dev=%d name=%s" % (i, name))
            break
if best is None:
    print("NO_NVIDIA")
    sys.exit(2)

nfam = ctypes.c_uint32(0)
vulkan.vkGetPhysicalDeviceQueueFamilyProperties(best, ctypes.byref(nfam),
                                                None)
fams = (VkQueueFamilyProperties * nfam.value)()
vulkan.vkGetPhysicalDeviceQueueFamilyProperties(best, ctypes.byref(nfam),
                                                fams)
gfam = None
for j in range(nfam.value):
    print("QFAM %d flags=0x%x count=%d" % (j, fams[j].queueFlags,
                                           fams[j].queueCount))
    if (fams[j].queueFlags & 0x1) and fams[j].queueCount > 0 and gfam is None:
        gfam = j
if gfam is None:
    print("NO_GRAPHICS_QUEUE")
    sys.exit(3)
prio = ctypes.c_float(1.0)
qci = VkDeviceQueueCreateInfo(VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
                              None, 0, gfam, 1, ctypes.pointer(prio))
dci = VkDeviceCreateInfo(VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, None, 0, 1,
                         ctypes.pointer(qci), 0, None, 0, None, None)
dev = ctypes.c_void_p()
r = vulkan.vkCreateDevice(best, ctypes.byref(dci), None, ctypes.byref(dev))
print("vkCreateDevice rc=%d (0x%x)" % (r, r if r < 0 else 0))
if r != 0:
    sys.exit(4)
queue = ctypes.c_void_p()
vulkan.vkGetDeviceQueue(dev, gfam, 0, ctypes.byref(queue))
print("QUEUE device=%s queue=%s" % (bool(dev), bool(queue)))
# optional: wait-idle (exercises the queue)
vulkan.vkDeviceWaitIdle(dev)
print("VK2_DEVICE_OK")
sys.exit(0)