typedef struct
{
    VkBuffer buffer;
    VkDeviceMemory memory;
}   vk_buffer;

typedef struct
{
    vk_buffer b;
    void *mapped;
}   vk_uniform;

static VkBuffer vk_buffer_create(vk_device *d, VkBufferUsageFlags usage, size_t size)
{
    VkBuffer buffer = {0};

    uint32_t queue_families[2] = {d->indices.graphics.index, d->indices.transfer.index};
    uint32_t queue_family_count = 2;
    bool separate_queues = (d->indices.graphics.index != d->indices.transfer.index) ? 1 : 0;

    VkBufferCreateInfo buffer_info = vk_buffer_info_init(separate_queues, queue_families, queue_family_count, usage, size);

    if (vkCreateBuffer(d->device, &buffer_info, NULL, &buffer) != VK_SUCCESS)
    {
        printf("Failed to create buffer\n");
    }

    return buffer;
}

static VkDeviceMemory vk_buffer_memory_init(vk_device *d, VkBuffer buffer, VkMemoryPropertyFlags properties)
{
    VkDeviceMemory memory = {0};

    VkMemoryRequirements mem_requirements = {0};
    vkGetBufferMemoryRequirements(d->device, buffer, &mem_requirements);

    VkMemoryAllocateInfo alloc_info = vk_memory_alloc_info_init(d, mem_requirements, properties);

    if (vkAllocateMemory(d->device, &alloc_info, NULL, &memory) != VK_SUCCESS)
    {
        printf("Failed to allocate vertex buffer memory\n");
    }

    if (vkBindBufferMemory(d->device, buffer, memory, 0) != VK_SUCCESS)
    {
        printf("Failed to bind buffer memory\n");
    }

    return memory;
}

static vk_buffer vk_buffer_init(vk_device *d, VkBufferUsageFlags usage, VkMemoryPropertyFlags properties, size_t size)
{
    vk_buffer b =
    {
        .buffer = vk_buffer_create(d, usage, size),
        .memory = vk_buffer_memory_init(d, b.buffer, properties)
    };

    return b;
}

vk_uniform vk_uniform_buffer_init(vk_device *d, size_t ub_size)
{
    vk_uniform ub = {0};
    ub.b = vk_buffer_init(d, VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, ub_size);

    vkMapMemory(d->device, ub.b.memory, 0, ub_size, 0, &ub.mapped);

    return ub;
}

vk_buffer vk_storage_buffer_init(vk_device *d)
{
    vk_buffer sb = vk_buffer_init(d, VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, sizeof(galaxy_particle) * PARTICLE_COUNT); // galaxy_particle = 32 bytes

    return sb;
}

typedef struct
{
    VkDescriptorSetLayout layout;
    VkDescriptorPool pool;
    VkDescriptorSet sets[MAX_FRAMES_IN_FLIGHT];
    vk_buffer particles[MAX_FRAMES_IN_FLIGHT];
    vk_uniform camera[MAX_FRAMES_IN_FLIGHT];
}   vk_desc_set;

VkDescriptorSetLayoutBinding vk_desc_set_layout_binding_init(uint32_t binding, VkDescriptorType type, uint32_t desc_count, VkShaderStageFlags stage)
{
    VkDescriptorSetLayoutBinding desc_set =
    {
        .binding = binding,
        .descriptorType = type,
        .descriptorCount = desc_count,
        .stageFlags = stage,
        .pImmutableSamplers = NULL
    };

    return desc_set;
}

static VkDescriptorSetLayoutCreateInfo vk_desc_set_layout_info_init(VkDescriptorSetLayoutBinding *binding, uint32_t binding_count)
{
    VkDescriptorSetLayoutCreateInfo layout_info =
    {
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        .pNext = NULL,
        .flags = 0,
        .bindingCount = binding_count,
        .pBindings = binding
    };

    return layout_info;
}

static VkDescriptorSetLayout vk_desc_set_layout_init(vk_device *d, VkDescriptorSetLayoutBinding *binding, uint32_t binding_count)
{
    VkDescriptorSetLayout layout = {0};

    VkDescriptorSetLayoutCreateInfo layout_info = vk_desc_set_layout_info_init(binding, binding_count);

    if (vkCreateDescriptorSetLayout(d->device, &layout_info, NULL, &layout) != VK_SUCCESS)
    {
        printf("Failed to create descriptor set layout\n");
    }

    return layout;
}

// Updates camera position; camera code not necessary for this problem
void vk_desc_set_update(vk_uniform *ub, camera *c, VkExtent2D img_extent)
{
    c->camera = camera_buffer_init(c, img_extent);
    memcpy(ub->mapped, &c->camera, sizeof(c->camera));
}

VkDescriptorPoolSize vk_desc_set_pool_size_init(VkDescriptorType type, uint32_t descriptor_count)
{
    VkDescriptorPoolSize size =
    {
        .type = type,
        .descriptorCount = descriptor_count
    };

    return size;
}

static VkDescriptorPoolCreateInfo vk_desc_set_pool_info_init(VkDescriptorPoolSize *size, uint32_t pool_count)
{
    VkDescriptorPoolCreateInfo pool_info =
    {
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
        .pNext = NULL,
        .flags = 0,
        .maxSets = MAX_FRAMES_IN_FLIGHT,
        .poolSizeCount = pool_count,
        .pPoolSizes = size
    };

    return pool_info;
}

static VkDescriptorPool vk_desc_set_pool_init(vk_device *d, VkDescriptorPoolSize *size, uint32_t size_count)
{
    VkDescriptorPool pool = {0};

    VkDescriptorPoolCreateInfo pool_info = vk_desc_set_pool_info_init(size, size_count);

    if (vkCreateDescriptorPool(d->device, &pool_info, NULL, &pool) != VK_SUCCESS)
    {
        printf("Failed to create descriptor pool\n");
    }

    return pool;
}

static VkDescriptorSetAllocateInfo vk_desc_set_alloc_info_init(VkDescriptorPool pool, VkDescriptorSetLayout *layouts)
{
    VkDescriptorSetAllocateInfo alloc_info =
    {
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
        .pNext = NULL,
        .descriptorPool = pool,
        .descriptorSetCount = MAX_FRAMES_IN_FLIGHT,
        .pSetLayouts = layouts
    };

    return alloc_info;
}

static VkDescriptorBufferInfo vk_desc_set_buffer_info_init(VkBuffer buffer, uint32_t offset, size_t size)
{
    VkDescriptorBufferInfo buffer_info =
    {
        .buffer = buffer,
        .offset = offset,
        .range = size
    };

    return buffer_info;
}

static VkWriteDescriptorSet vk_desc_set_write_info_init(VkDescriptorSet set, uint32_t binding, VkDescriptorType type, VkDescriptorBufferInfo *buffer_info)
{
    VkWriteDescriptorSet write_info =
    {
        .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .pNext = NULL,
        .dstSet = set,
        .dstBinding = binding,
        .dstArrayElement = 0,
        .descriptorCount = 1,
        .descriptorType = type,
        .pImageInfo = NULL,
        .pBufferInfo = buffer_info,
        .pTexelBufferView = NULL
    };

    return write_info;
}

void vk_desc_galaxy_sets_init(vk_device *d, vk_desc_set *ds)
{
    VkDescriptorSetLayout layouts[MAX_FRAMES_IN_FLIGHT] = {0};

    for (uint32_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
    {
        layouts[i] = ds->layout;
        ds->camera[i] = vk_uniform_buffer_init(d, sizeof(mat4s));
        ds->particles[i] = vk_storage_buffer_init(d);
    }

    VkDescriptorSetAllocateInfo alloc_info = vk_desc_set_alloc_info_init(ds->pool, layouts);

    if (vkAllocateDescriptorSets(d->device, &alloc_info, ds->sets) != VK_SUCCESS)
    {
        printf("Failed to allocate galaxy descriptor sets\n");
    }

    for (uint32_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
    {
        VkDescriptorBufferInfo uniform_info = vk_desc_set_buffer_info_init(ds->camera[i].b.buffer, 0, sizeof(mat4s));
        VkWriteDescriptorSet uniform_write_info = vk_desc_set_write_info_init(ds->sets[i], 0, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, &uniform_info);

        VkDescriptorBufferInfo storage_info = vk_desc_set_buffer_info_init(ds->particles[i].buffer, 0, VK_WHOLE_SIZE);
        VkWriteDescriptorSet storage_write_info = vk_desc_set_write_info_init(ds->sets[i], 1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC, &storage_info);

        VkWriteDescriptorSet desc_writes[2] = {uniform_write_info, storage_write_info};

        vkUpdateDescriptorSets(d->device, 2, desc_writes, 0, NULL);
    }
}

void vk_desc_compute_sets_init(vk_device *d, vk_desc_set *ds)
{
    VkDescriptorSetLayout layouts[MAX_FRAMES_IN_FLIGHT] = {0};

    for (uint32_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
    {
        layouts[i] = ds->layout;
        ds->particles[i] = vk_storage_buffer_init(d);
    }

    VkDescriptorSetAllocateInfo alloc_info = vk_desc_set_alloc_info_init(ds->pool, layouts);

    if (vkAllocateDescriptorSets(d->device, &alloc_info, ds->sets) != VK_SUCCESS)
    {
        printf("Failed to allocate compute descriptor sets\n");
    }

    for (uint32_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
    {
        VkDescriptorBufferInfo storage_info = vk_desc_set_buffer_info_init(ds->particles[i].buffer, 0, VK_WHOLE_SIZE);
        VkWriteDescriptorSet storage_write_info = vk_desc_set_write_info_init(ds->sets[i], 0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC, &storage_info);

        vkUpdateDescriptorSets(d->device, 1, &storage_write_info, 0, NULL);
    }
}

vk_desc_set vk_desc_set_init(vk_device *d, VkDescriptorSetLayoutBinding *binding, uint32_t binding_count, VkDescriptorPoolSize *size, uint32_t size_count)
{
    vk_desc_set ds =
    {
        .layout = vk_desc_set_layout_init(d, binding, binding_count),
        .pool = vk_desc_set_pool_init(d, size, size_count)
    };

    return ds;
}