#!/usr/bin/env python3
# vk_alloc_test.py - reproduce vk_frame's DEVICE_LOCAL image allocation:
# instance -> device -> OPTIMAL R8G8B8A8 storage image -> mem reqs ->
# pick first DEVICE_LOCAL type -> vkAllocateMemory -> report rc/type/size.
import ctypes, sys

vk = ctypes.CDLL("libvulkan.so.1")

class VkApplicationInfo(ctypes.Structure):
    _fields_ = [("sType", ctypes.c_uint32), ("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_uint32), ("pNext", ctypes.c_void_p),
                ("flags", ctypes.c_uint32), ("pApplicationInfo", ctypes.c_void_p),
                ("enabledLayerCount", ctypes.c_uint32), ("ppEnabledLayerNames", ctypes.c_void_p),
                ("enabledExtensionCount", ctypes.c_uint32), ("ppEnabledExtensionNames", ctypes.c_void_p)]

class VkMemoryType(ctypes.Structure):
    _fields_ = [("propertyFlags", ctypes.c_uint32), ("heapIndex", ctypes.c_uint32)]
class VkMemoryHeap(ctypes.Structure):
    _fields_ = [("size", ctypes.c_uint64), ("flags", ctypes.c_uint32)]
class VkPhysicalDeviceMemoryProperties(ctypes.Structure):
    _fields_ = [("memoryTypeCount", ctypes.c_uint32),
                ("memoryTypes", VkMemoryType * 32),
                ("memoryHeapCount", ctypes.c_uint32),
                ("memoryHeaps", VkMemoryHeap * 16)]

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

class VkExtent3D(ctypes.Structure):
    _fields_ = [("width", ctypes.c_uint32), ("height", ctypes.c_uint32), ("depth", ctypes.c_uint32)]
class VkImageCreateInfo(ctypes.Structure):
    _fields_ = [("sType", ctypes.c_uint32), ("pNext", ctypes.c_void_p), ("flags", ctypes.c_uint32),
                ("imageType", ctypes.c_uint32), ("format", ctypes.c_uint32),
                ("extent", VkExtent3D), ("mipLevels", ctypes.c_uint32), ("arrayLayers", ctypes.c_uint32),
                ("samples", ctypes.c_uint32), ("tiling", ctypes.c_uint32), ("usage", ctypes.c_uint32),
                ("sharingMode", ctypes.c_uint32), ("queueFamilyIndexCount", ctypes.c_uint32),
                ("pQueueFamilyIndices", ctypes.c_void_p), ("initialLayout", ctypes.c_uint32)]
class VkMemoryRequirements(ctypes.Structure):
    _fields_ = [("size", ctypes.c_uint64), ("alignment", ctypes.c_uint64), ("memoryTypeBits", ctypes.c_uint32)]
class VkMemoryAllocateInfo(ctypes.Structure):
    _fields_ = [("sType", ctypes.c_uint32), ("pNext", ctypes.c_void_p),
                ("allocationSize", ctypes.c_uint64), ("memoryTypeIndex", ctypes.c_uint32)]

app = VkApplicationInfo(0, None, b"alloctest", 1, b"alloctest", 1, 0x401000)
ci = VkInstanceCreateInfo(1, None, 0, ctypes.cast(ctypes.pointer(app), ctypes.c_void_p), 0, None, 0, None)
inst = ctypes.c_void_p()
print("vkCreateInstance rc=%d" % vk.vkCreateInstance(ctypes.byref(ci), None, ctypes.byref(inst)))

n = ctypes.c_uint32(0)
vk.vkEnumeratePhysicalDevices(inst, ctypes.byref(n), None)
devs = (ctypes.c_void_p * n.value)()
vk.vkEnumeratePhysicalDevices(inst, ctypes.byref(n), devs)
phys = devs[0]
print("phys=0x%x" % (phys or 0))

prio = ctypes.c_float(1.0)
qci = VkDeviceQueueCreateInfo(2, None, 0, 0, 1, ctypes.cast(ctypes.pointer(prio), ctypes.c_void_p))
dci = VkDeviceCreateInfo(3, None, 0, 1, ctypes.cast(ctypes.pointer(qci), ctypes.c_void_p), 0, None, 0, None, None)
dev = ctypes.c_void_p()
rc = vk.vkCreateDevice(phys, ctypes.byref(dci), None, ctypes.byref(dev))
print("vkCreateDevice rc=%d dev=0x%x" % (rc, dev or 0))
if rc != 0:
    sys.exit(1)

mp = VkPhysicalDeviceMemoryProperties()
vk.vkGetPhysicalDeviceMemoryProperties(phys, ctypes.byref(mp))

img_ci = VkImageCreateInfo(14, None, 0, 1, 37,  # VK_IMAGE_TYPE_2D, R8G8B8A8_UNORM
                           VkExtent3D(64, 64, 1), 1, 1, 1, 0,  # OPTIMAL
                           0x2 | 0x1, 0, 0, None, 0)  # STORAGE | TRANSFER_SRC, UNDEFINED
img = ctypes.c_void_p()
rc = vk.vkCreateImage(dev, ctypes.byref(img_ci), None, ctypes.byref(img))
print("vkCreateImage rc=%d img=0x%x" % (rc, img or 0))
if rc != 0:
    sys.exit(1)

mr = VkMemoryRequirements()
vk.vkGetImageMemoryRequirements(dev, img, ctypes.byref(mr))
print("memreqs size=0x%x align=0x%x typeBits=0x%x" % (mr.size, mr.alignment, mr.memoryTypeBits))

tidx = 0
found = False
for k in range(mp.memoryTypeCount):
    if (mr.memoryTypeBits & (1 << k)) and (mp.memoryTypes[k].propertyFlags & 0x1):
        tidx = k
        found = True
        break
print("selected type=%d found_device_local=%s (types: %s)" %
      (tidx, found, ", ".join("%d:0x%x" % (k, mp.memoryTypes[k].propertyFlags) for k in range(mp.memoryTypeCount))))

ma = VkMemoryAllocateInfo(5, None, mr.size, tidx)
mem = ctypes.c_void_p()
rc = vk.vkAllocateMemory(dev, ctypes.byref(ma), None, ctypes.byref(mem))
print("vkAllocateMemory size=0x%x type=%d rc=%d" % (mr.size, tidx, rc))