#if 0 /*
#/  ================================================================
#/
#/    vulkan_compute_minimal_example.c
#/
#/  ----------------------------------------------------------------
#/
#/    (C) 2025 Mitya Selivanov <guattari.tech>
#/
#/  ================================================================
#/
#/  COMPUTE SHADER SOURCE CODE
#/
#/  ================================================================
cat <<EOF >~module.comp

#version 450
layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
layout(binding = 0) buffer output_buffer {
  uint values_out[16];
};
void main() {
  for (uint i = 0; i < 16; i++) values_out[i] = i;
}

EOF
#/  ================================================================
#/
#/  BUILD AND RUN SCRIPT
#/
#/  ================================================================

SRC=${0##*./}
BIN=${SRC%.*}
glslc ~module.comp -o ~module.spv && gcc $SRC -o $BIN -lm -lvulkan && ./$BIN
STATUS=$?
rm -f $BIN
rm -f ~module.comp
rm -f ~module.spv
exit $STATUS # */
#endif

//  ================================================================
//
//  C SOURCE CODE
//
//  ================================================================

#include <string.h>
#include <stdio.h>

#include <vulkan/vulkan.h>

typedef signed   i32;
typedef unsigned u32;
typedef float    f32;
typedef char     c8;

enum {
  MAX_DEVICE_COUNT       = 16,
  MAX_QUEUE_FAMILY_COUNT = 64,
  BUFFER_SIZE            = 64,
  GROUP_COUNT_X          = 1,
  GROUP_COUNT_Y          = 1,
  GROUP_COUNT_Z          = 1,
};

i32 main(i32 argc, c8 **argv) {
  i32 status             = 0;
  u32 queue_family_index = 0;

  VkInstance            instance              = NULL;
  VkPhysicalDevice      physical_device       = NULL;
  VkDevice              device                = NULL;
  VkQueue               queue                 = NULL;
  VkBuffer              buffer                = NULL;
  VkDeviceMemory        buffer_memory         = NULL;
  VkDescriptorSetLayout descriptor_set_layout = NULL;
  VkDescriptorPool      descriptor_pool       = NULL;
  VkDescriptorSet       descriptor_set        = NULL;
  VkShaderModule        compute_module        = NULL;
  VkPipelineLayout      pipeline_layout       = NULL;
  VkPipeline            pipeline              = NULL;
  VkCommandPool         command_pool          = NULL;
  VkCommandBuffer       command_buffer        = NULL;
  VkFence               fence                 = NULL;

  if (vkCreateInstance(&(VkInstanceCreateInfo) {
    .sType            = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
    .pApplicationInfo = &(VkApplicationInfo) {
      .sType      = VK_STRUCTURE_TYPE_APPLICATION_INFO,
      .apiVersion = VK_API_VERSION_1_1,
    },
  }, NULL, &instance) != VK_SUCCESS) {
    printf("vkCreateInstance failed.\n");
    status = -1;
    goto _cleanup;
  }

  //  Find physical device
  {
    u32              device_count = MAX_DEVICE_COUNT;
    VkPhysicalDevice devices[MAX_DEVICE_COUNT];

    VkResult res = vkEnumeratePhysicalDevices(instance, &device_count, devices);

    if (res != VK_SUCCESS && res != VK_INCOMPLETE) {
      printf("vkEnumeratePhysicalDevices failed.\n");
      status = -1;
      goto _cleanup;
    }

    if (device_count <= 0) {
      printf("Compatible physical device not found.\n");
      status = -1;
      goto _cleanup;
    }

    u32                        i = 0;
    VkPhysicalDeviceProperties properties;

    for (; i < device_count; ++i) {
      vkGetPhysicalDeviceProperties(devices[i], &properties);

      if (strstr(properties.deviceName, "NVIDIA") != NULL || strstr(properties.deviceName, "AMD") != NULL) {
        physical_device = devices[i];
        break;
      }
    }

    if (i >= device_count)
      physical_device = devices[0];

    vkGetPhysicalDeviceProperties(physical_device, &properties);

    printf("Physical device selected: %s\n", properties.deviceName);
  }

  //  Find queue family
  {
    u32                     queue_family_count = MAX_QUEUE_FAMILY_COUNT;
    VkQueueFamilyProperties queue_families[MAX_QUEUE_FAMILY_COUNT];

    vkGetPhysicalDeviceQueueFamilyProperties(physical_device, &queue_family_count, queue_families);

    u32 index = 0;

    for (; index < queue_family_count; ++index)
      if (queue_families[index].queueCount > 0 && (queue_families[index].queueFlags & VK_QUEUE_COMPUTE_BIT))
        break;

    if (index >= queue_family_count) {
      printf("Compatible queue family not found.\n");
      status = -1;
      goto _cleanup;
    }

    queue_family_index = index;
  }

  //  Create logical device
  {
    VkPhysicalDeviceFeatures device_features;
    memset(&device_features, 0, sizeof device_features);

    if (vkCreateDevice(physical_device, &(VkDeviceCreateInfo) {
      .sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
      .pQueueCreateInfos = &(VkDeviceQueueCreateInfo) {
        .sType            = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
        .queueCount       = 1,
        .pQueuePriorities = (f32[1]) { 1.f, },
        .queueFamilyIndex = queue_family_index,
      },
      .queueCreateInfoCount = 1,
      .pEnabledFeatures = &device_features,
    }, NULL, &device) != VK_SUCCESS) {
      printf("vkCreateDevice failed.\n");
      status = -1;
      goto _cleanup;
    }

    vkGetDeviceQueue(device, queue_family_index, 0, &queue);
  }

  //  Create buffer
  {
    if (vkCreateBuffer(device, &(VkBufferCreateInfo) {
        .sType       = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
        .size        = BUFFER_SIZE,
        .usage       = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
        .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
    }, NULL, &buffer) != VK_SUCCESS) {
      printf("vkCreateBuffer failed.\n");
      status = -1;
      goto _cleanup;
    }

    VkMemoryRequirements             memory_requirements;
    VkPhysicalDeviceMemoryProperties memory_properties;

    vkGetBufferMemoryRequirements(device, buffer, &memory_requirements);
    vkGetPhysicalDeviceMemoryProperties(physical_device, &memory_properties);

    u32 i                 = 0;
    u32 properties        = VK_MEMORY_PROPERTY_HOST_COHERENT_BIT | VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT;
    u32 memory_type_index = 0;

    for (; i < memory_properties.memoryTypeCount; ++i)
      if ((memory_requirements.memoryTypeBits & (1 << i)) && ((memory_properties.memoryTypes[i].propertyFlags & properties) == properties)) {
        memory_type_index = i;
        break;
      }

    if (i >= memory_properties.memoryTypeCount) {
      printf("Compatible memory properties not found.\n");
      status = -1;
      goto _cleanup;
    }

    if (vkAllocateMemory(device, &(VkMemoryAllocateInfo) {
      .sType           = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
      .allocationSize  = memory_requirements.size,
      .memoryTypeIndex = memory_type_index,
    }, NULL, &buffer_memory) != VK_SUCCESS) {
      printf("vkAllocateMemory failed.\n");
      status = -1;
      goto _cleanup;
    }

    if (vkBindBufferMemory(device, buffer, buffer_memory, 0) != VK_SUCCESS) {
      printf("vkBindBufferMemory failed.\n");
      status = -1;
      goto _cleanup;
    }
  }

  if (vkCreateDescriptorSetLayout(device, &(VkDescriptorSetLayoutCreateInfo) {
    .sType        = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
    .bindingCount = 1,
    .pBindings    = &(VkDescriptorSetLayoutBinding) {
      .binding         = 0,
      .descriptorType  = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
      .descriptorCount = 1,
      .stageFlags      = VK_SHADER_STAGE_COMPUTE_BIT,
    },
  }, NULL, &descriptor_set_layout) != VK_SUCCESS) {
    printf("vkCreateDescriptorSetLayout failed.\n");
    status = -1;
    goto _cleanup;
  }

  //  Create descriptor set
  {
    if (vkCreateDescriptorPool(device, &(VkDescriptorPoolCreateInfo) {
      .sType         = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
      .maxSets       = 1,
      .poolSizeCount = 1,
      .pPoolSizes    = &(VkDescriptorPoolSize) {
        .type            = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .descriptorCount = 1,
      },
    }, NULL, &descriptor_pool) != VK_SUCCESS) {
      printf("vkCreateDescriptorPool failed.\n");
      status = -1;
      goto _cleanup;
    }

    if (vkAllocateDescriptorSets(device, &(VkDescriptorSetAllocateInfo) {
      .sType              = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
      .descriptorPool     = descriptor_pool,
      .descriptorSetCount = 1,
      .pSetLayouts        = &descriptor_set_layout,
    }, &descriptor_set) != VK_SUCCESS) {
      printf("vkAllocateDescriptorSets failed.\n");
      status = -1;
      goto _cleanup;
    }

    vkUpdateDescriptorSets(device, 1, &(VkWriteDescriptorSet) {
      .sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
      .dstSet          = descriptor_set,
      .dstBinding      = 0,
      .descriptorCount = 1,
      .descriptorType  = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
      .pBufferInfo     = &(VkDescriptorBufferInfo) {
        .buffer = buffer,
        .offset = 0,
        .range  = BUFFER_SIZE,
      },
    }, 0, NULL);
  }

  //  Create compute pipeline
  {
    FILE *f = fopen("~module.spv", "rb");
    if (f == NULL) {
      printf("fopen failed.\n");
      status = -1;
      goto _cleanup;
    }
    static u32 code[4096];
    u32 n = fread(code, 1, sizeof code, f);
    fclose(f);

    if (vkCreateShaderModule(device, &(VkShaderModuleCreateInfo) {
      .sType    = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
      .pCode    = code,
      .codeSize = n,
    }, NULL, &compute_module) != VK_SUCCESS) {
      printf("vkCreateShaderModule failed.\n");
      status = -1;
      goto _cleanup;
    }

    if (vkCreatePipelineLayout(device, &(VkPipelineLayoutCreateInfo) {
      .sType          = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
      .setLayoutCount = 1,
      .pSetLayouts    = &descriptor_set_layout,
    }, NULL, &pipeline_layout) != VK_SUCCESS) {
      printf("vkCreatePipelineLayout failed.\n");
      status = -1;
      goto _cleanup;
    }

    if (vkCreateComputePipelines(device, VK_NULL_HANDLE, 1, &(VkComputePipelineCreateInfo) {
      .sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
      .stage = {
        .sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
        .stage  = VK_SHADER_STAGE_COMPUTE_BIT,
        .module = compute_module,
        .pName  = "main",
      },
      .layout = pipeline_layout,
    }, NULL, &pipeline) != VK_SUCCESS) {
      printf("vkCreateComputePipelines failed.\n");
      status = -1;
      goto _cleanup;
    }
  }

  //  Create command buffer
  {
    if (vkCreateCommandPool(device, &(VkCommandPoolCreateInfo) {
      .sType            = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
      .queueFamilyIndex = queue_family_index,
    }, NULL, &command_pool) != VK_SUCCESS) {
      printf("vkCreateCommandPool failed.\n");
      status = -1;
      goto _cleanup;
    }

    if (vkAllocateCommandBuffers(device, &(VkCommandBufferAllocateInfo) {
      .sType              = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
      .commandPool        = command_pool,
      .level              = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
      .commandBufferCount = 1,
    }, &command_buffer) != VK_SUCCESS) {
      printf("vkAllocateCommandBuffers failed.\n");
      status = -1;
      goto _cleanup;
    }

    if (vkBeginCommandBuffer(command_buffer, &(VkCommandBufferBeginInfo) {
      .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
      .flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT,
    }) != VK_SUCCESS) {
      printf("vkBeginCommandBuffer failed.\n");
      status = -1;
      goto _cleanup;
    }

    vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline);
    vkCmdBindDescriptorSets(command_buffer, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline_layout, 0, 1, &descriptor_set, 0, NULL);
    vkCmdDispatch(command_buffer, GROUP_COUNT_X, GROUP_COUNT_Y, GROUP_COUNT_Z);

    if (vkEndCommandBuffer(command_buffer) != VK_SUCCESS) {
      printf("vkEndCommandBuffer failed.\n");
      status = -1;
      goto _cleanup;
    }
  }

  //  Run command buffer
  {
    if (vkCreateFence(device, &(VkFenceCreateInfo) { .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO, }, NULL, &fence) != VK_SUCCESS) {
      printf("vkCreateFence failed.\n");
      status = -1;
      goto _cleanup;
    }

    if (vkQueueSubmit(queue, 1, &(VkSubmitInfo) {
      .sType              = VK_STRUCTURE_TYPE_SUBMIT_INFO,
      .commandBufferCount = 1,
      .pCommandBuffers    = &command_buffer,
    }, fence) != VK_SUCCESS) {
      printf("vkQueueSubmit failed.\n");
      status = -1;
      goto _cleanup;
    }

    if (vkWaitForFences(device, 1, &fence, VK_TRUE, 100000000000) != VK_SUCCESS) {
      printf("vkWaitForFences failed.\n");
      status = -1;
      goto _cleanup;
    }
  }

  //  Read output buffer
  {
    u32 *p = NULL;

    if (vkMapMemory(device, buffer_memory, 0, BUFFER_SIZE, 0, (void **) &p) != VK_SUCCESS) {
      printf("vkMapMemory failed.\n");
      status = -1;
      goto _cleanup;
    }

    printf("Data:");
    for (u32 i = 0; i < BUFFER_SIZE / 4; i++)
      printf(" %2d", (int) p[i]);
    printf("\n");

    vkUnmapMemory(device, buffer_memory);
  }

_cleanup:
  {
    if (fence != NULL)
      vkDestroyFence(device, fence, NULL);
    if (buffer_memory != NULL)
      vkFreeMemory(device, buffer_memory, NULL);
    if (buffer != NULL)
      vkDestroyBuffer(device, buffer, NULL);
    if (compute_module != NULL)
      vkDestroyShaderModule(device, compute_module, NULL);  
    if (descriptor_pool != NULL)
      vkDestroyDescriptorPool(device, descriptor_pool, NULL);
    if (descriptor_set_layout != NULL)
      vkDestroyDescriptorSetLayout(device, descriptor_set_layout, NULL);
    if (pipeline_layout != NULL)
      vkDestroyPipelineLayout(device, pipeline_layout, NULL);
    if (pipeline != NULL)
      vkDestroyPipeline(device, pipeline, NULL);
    if (command_pool != NULL)
      vkDestroyCommandPool(device, command_pool, NULL);
    if (device != NULL)
      vkDestroyDevice(device, NULL);
    if (instance != NULL)
      vkDestroyInstance(instance, NULL);

    instance              = NULL;
    physical_device       = NULL;
    device                = NULL;
    queue                 = NULL;
    buffer                = NULL;
    buffer_memory         = NULL;
    descriptor_set_layout = NULL;
    descriptor_pool       = NULL;
    descriptor_set        = NULL;
    compute_module        = NULL;
    pipeline_layout       = NULL;
    pipeline              = NULL;
    command_pool          = NULL;
    command_buffer        = NULL;
    fence                 = NULL;
  }

  printf("OK\n");
  return status;
}