package pipeline

import gpu "../core"
import "core:log"
import vk "vendor:vulkan"

Descriptor_State :: struct {
	set_layout: vk.DescriptorSetLayout,
	pool:       vk.DescriptorPool,
}

create_descriptor_set_layout :: proc(
	dev: ^gpu.Gpu_Device,
) -> (
	layout: vk.DescriptorSetLayout,
	ok: bool,
) {
	binding := vk.DescriptorSetLayoutBinding {
		binding         = 0,
		descriptorType  = .COMBINED_IMAGE_SAMPLER,
		descriptorCount = 1,
		stageFlags      = {.FRAGMENT},
	}

	create_info := vk.DescriptorSetLayoutCreateInfo {
		sType        = .DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
		bindingCount = 1,
		pBindings    = &binding,
	}

	result := vk.CreateDescriptorSetLayout(dev.device, &create_info, nil, &layout)
	if result != .SUCCESS {
		log.errorf("gpu/pipeline: vkCreateDescriptorSetLayout failed: %v", result)
		return {}, false
	}

	return layout, true
}

create_descriptor_pool :: proc(
	dev: ^gpu.Gpu_Device,
	max_sets: u32,
) -> (
	pool: vk.DescriptorPool,
	ok: bool,
) {
	pool_size := vk.DescriptorPoolSize {
		type            = .COMBINED_IMAGE_SAMPLER,
		descriptorCount = max_sets,
	}

	create_info := vk.DescriptorPoolCreateInfo {
		sType         = .DESCRIPTOR_POOL_CREATE_INFO,
		flags         = {.FREE_DESCRIPTOR_SET},
		maxSets       = max_sets,
		poolSizeCount = 1,
		pPoolSizes    = &pool_size,
	}

	result := vk.CreateDescriptorPool(dev.device, &create_info, nil, &pool)
	if result != .SUCCESS {
		log.errorf("gpu/pipeline: vkCreateDescriptorPool failed: %v", result)
		return {}, false
	}

	return pool, true
}

allocate_descriptor_set :: proc(
	dev: ^gpu.Gpu_Device,
	pool: vk.DescriptorPool,
	layout: vk.DescriptorSetLayout,
) -> (
	set: vk.DescriptorSet,
	ok: bool,
) {
	layout_copy := layout
	alloc_info := vk.DescriptorSetAllocateInfo {
		sType              = .DESCRIPTOR_SET_ALLOCATE_INFO,
		descriptorPool     = pool,
		descriptorSetCount = 1,
		pSetLayouts        = &layout_copy,
	}

	result := vk.AllocateDescriptorSets(dev.device, &alloc_info, &set)
	if result != .SUCCESS {
		log.errorf("gpu/pipeline: vkAllocateDescriptorSets failed: %v", result)
		return {}, false
	}

	return set, true
}

update_descriptor_set_image :: proc(
	dev: ^gpu.Gpu_Device,
	set: vk.DescriptorSet,
	image_view: vk.ImageView,
	sampler: vk.Sampler,
	image_layout: vk.ImageLayout = .SHADER_READ_ONLY_OPTIMAL,
) {
	image_info := vk.DescriptorImageInfo {
		sampler     = sampler,
		imageView   = image_view,
		imageLayout = image_layout,
	}

	write := vk.WriteDescriptorSet {
		sType           = .WRITE_DESCRIPTOR_SET,
		dstSet          = set,
		dstBinding      = 0,
		dstArrayElement = 0,
		descriptorCount = 1,
		descriptorType  = .COMBINED_IMAGE_SAMPLER,
		pImageInfo      = &image_info,
	}

	vk.UpdateDescriptorSets(dev.device, 1, &write, 0, nil)
}

// Create a descriptor set layout for a uniform buffer (e.g. light UBO)
create_ubo_descriptor_set_layout :: proc(
	dev: ^gpu.Gpu_Device,
	stage_flags: vk.ShaderStageFlags = {.VERTEX, .FRAGMENT},
) -> (
	layout: vk.DescriptorSetLayout,
	ok: bool,
) {
	binding := vk.DescriptorSetLayoutBinding {
		binding         = 0,
		descriptorType  = .UNIFORM_BUFFER,
		descriptorCount = 1,
		stageFlags      = stage_flags,
	}

	create_info := vk.DescriptorSetLayoutCreateInfo {
		sType        = .DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
		bindingCount = 1,
		pBindings    = &binding,
	}

	result := vk.CreateDescriptorSetLayout(dev.device, &create_info, nil, &layout)
	if result != .SUCCESS {
		log.errorf("gpu/pipeline: vkCreateDescriptorSetLayout (UBO) failed: %v", result)
		return {}, false
	}

	return layout, true
}

// Create a descriptor pool for uniform buffers
create_ubo_descriptor_pool :: proc(
	dev: ^gpu.Gpu_Device,
	max_sets: u32,
) -> (
	pool: vk.DescriptorPool,
	ok: bool,
) {
	pool_size := vk.DescriptorPoolSize {
		type            = .UNIFORM_BUFFER,
		descriptorCount = max_sets,
	}

	create_info := vk.DescriptorPoolCreateInfo {
		sType         = .DESCRIPTOR_POOL_CREATE_INFO,
		flags         = {.FREE_DESCRIPTOR_SET},
		maxSets       = max_sets,
		poolSizeCount = 1,
		pPoolSizes    = &pool_size,
	}

	result := vk.CreateDescriptorPool(dev.device, &create_info, nil, &pool)
	if result != .SUCCESS {
		log.errorf("gpu/pipeline: vkCreateDescriptorPool (UBO) failed: %v", result)
		return {}, false
	}

	return pool, true
}

// Write a buffer binding to a descriptor set
update_descriptor_set_buffer :: proc(
	dev: ^gpu.Gpu_Device,
	set: vk.DescriptorSet,
	buffer: vk.Buffer,
	size: vk.DeviceSize,
	binding: u32 = 0,
) {
	buffer_info := vk.DescriptorBufferInfo {
		buffer = buffer,
		offset = 0,
		range  = size,
	}

	write := vk.WriteDescriptorSet {
		sType           = .WRITE_DESCRIPTOR_SET,
		dstSet          = set,
		dstBinding      = binding,
		dstArrayElement = 0,
		descriptorCount = 1,
		descriptorType  = .UNIFORM_BUFFER,
		pBufferInfo     = &buffer_info,
	}

	vk.UpdateDescriptorSets(dev.device, 1, &write, 0, nil)
}

destroy_descriptor_pool :: proc(dev: ^gpu.Gpu_Device, pool: ^vk.DescriptorPool) {
	if pool^ != 0 {
		vk.DestroyDescriptorPool(dev.device, pool^, nil)
		pool^ = 0
	}
}

destroy_descriptor_set_layout :: proc(dev: ^gpu.Gpu_Device, layout: ^vk.DescriptorSetLayout) {
	if layout^ != 0 {
		vk.DestroyDescriptorSetLayout(dev.device, layout^, nil)
		layout^ = 0
	}
}