package pipeline

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

Compute_Pipeline :: struct {
	pipeline:        vk.Pipeline,
	layout:          vk.PipelineLayout,
	desc_set_layout: vk.DescriptorSetLayout,
	desc_pool:       vk.DescriptorPool,
	num_buffers:     u32,
}

create_compute_pipeline :: proc(
	dev: ^gpu.Gpu_Device,
	shader_module: vk.ShaderModule,
	num_storage_buffers: u32,
	push_constant_size: u32 = 0,
) -> (
	cp: Compute_Pipeline,
	ok: bool,
) {
	cp.num_buffers = num_storage_buffers

	// Create descriptor set layout with sequential STORAGE_BUFFER bindings
	if num_storage_buffers > 0 {
		bindings := make(
			[]vk.DescriptorSetLayoutBinding,
			num_storage_buffers,
			context.temp_allocator,
		)
		for i in 0 ..< num_storage_buffers {
			bindings[i] = vk.DescriptorSetLayoutBinding {
				binding         = i,
				descriptorType  = .STORAGE_BUFFER,
				descriptorCount = 1,
				stageFlags      = {.COMPUTE},
			}
		}

		layout_info := vk.DescriptorSetLayoutCreateInfo {
			sType        = .DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
			bindingCount = num_storage_buffers,
			pBindings    = raw_data(bindings),
		}

		result := vk.CreateDescriptorSetLayout(dev.device, &layout_info, nil, &cp.desc_set_layout)
		if result != .SUCCESS {
			log.errorf("gpu/pipeline: compute descriptor set layout failed: %v", result)
			return {}, false
		}

		// Create descriptor pool for MAX_FRAMES_IN_FLIGHT sets
		pool_size := vk.DescriptorPoolSize {
			type            = .STORAGE_BUFFER,
			descriptorCount = num_storage_buffers * gpu.MAX_FRAMES_IN_FLIGHT,
		}

		pool_info := vk.DescriptorPoolCreateInfo {
			sType         = .DESCRIPTOR_POOL_CREATE_INFO,
			flags         = {.FREE_DESCRIPTOR_SET},
			maxSets       = gpu.MAX_FRAMES_IN_FLIGHT,
			poolSizeCount = 1,
			pPoolSizes    = &pool_size,
		}

		result = vk.CreateDescriptorPool(dev.device, &pool_info, nil, &cp.desc_pool)
		if result != .SUCCESS {
			log.errorf("gpu/pipeline: compute descriptor pool failed: %v", result)
			vk.DestroyDescriptorSetLayout(dev.device, cp.desc_set_layout, nil)
			return {}, false
		}
	}

	// Pipeline layout
	push_range := vk.PushConstantRange {
		stageFlags = {.COMPUTE},
		offset     = 0,
		size       = push_constant_size,
	}

	layout_info := vk.PipelineLayoutCreateInfo {
		sType = .PIPELINE_LAYOUT_CREATE_INFO,
	}
	if num_storage_buffers > 0 {
		layout_info.setLayoutCount = 1
		layout_info.pSetLayouts = &cp.desc_set_layout
	}
	if push_constant_size > 0 {
		layout_info.pushConstantRangeCount = 1
		layout_info.pPushConstantRanges = &push_range
	}

	result := vk.CreatePipelineLayout(dev.device, &layout_info, nil, &cp.layout)
	if result != .SUCCESS {
		log.errorf("gpu/pipeline: compute pipeline layout failed: %v", result)
		destroy_compute_pipeline(dev, &cp)
		return {}, false
	}

	// Compute pipeline
	stage := vk.PipelineShaderStageCreateInfo {
		sType  = .PIPELINE_SHADER_STAGE_CREATE_INFO,
		stage  = {.COMPUTE},
		module = shader_module,
		pName  = "main",
	}

	pipeline_info := vk.ComputePipelineCreateInfo {
		sType  = .COMPUTE_PIPELINE_CREATE_INFO,
		stage  = stage,
		layout = cp.layout,
	}

	result = vk.CreateComputePipelines(dev.device, 0, 1, &pipeline_info, nil, &cp.pipeline)
	if result != .SUCCESS {
		log.errorf("gpu/pipeline: vkCreateComputePipelines failed: %v", result)
		destroy_compute_pipeline(dev, &cp)
		return {}, false
	}

	return cp, true
}

destroy_compute_pipeline :: proc(dev: ^gpu.Gpu_Device, cp: ^Compute_Pipeline) {
	if cp.pipeline != 0 {
		vk.DestroyPipeline(dev.device, cp.pipeline, nil)
		cp.pipeline = 0
	}
	if cp.layout != 0 {
		vk.DestroyPipelineLayout(dev.device, cp.layout, nil)
		cp.layout = 0
	}
	if cp.desc_pool != 0 {
		vk.DestroyDescriptorPool(dev.device, cp.desc_pool, nil)
		cp.desc_pool = 0
	}
	if cp.desc_set_layout != 0 {
		vk.DestroyDescriptorSetLayout(dev.device, cp.desc_set_layout, nil)
		cp.desc_set_layout = 0
	}
}

allocate_compute_descriptor_sets :: proc(
	dev: ^gpu.Gpu_Device,
	cp: ^Compute_Pipeline,
	count: u32,
) -> (
	sets: []vk.DescriptorSet,
	ok: bool,
) {
	if cp.desc_pool == 0 || cp.desc_set_layout == 0 {
		return nil, false
	}

	layouts := make([]vk.DescriptorSetLayout, count, context.temp_allocator)
	for i in 0 ..< count {
		layouts[i] = cp.desc_set_layout
	}

	sets = make([]vk.DescriptorSet, count)
	alloc_info := vk.DescriptorSetAllocateInfo {
		sType              = .DESCRIPTOR_SET_ALLOCATE_INFO,
		descriptorPool     = cp.desc_pool,
		descriptorSetCount = count,
		pSetLayouts        = raw_data(layouts),
	}

	result := vk.AllocateDescriptorSets(dev.device, &alloc_info, raw_data(sets))
	if result != .SUCCESS {
		log.errorf("gpu/pipeline: compute descriptor set allocation failed: %v", result)
		delete(sets)
		return nil, false
	}

	return sets, true
}

update_compute_descriptor_buffer :: proc(
	dev: ^gpu.Gpu_Device,
	set: vk.DescriptorSet,
	binding: u32,
	buffer: vk.Buffer,
	size: vk.DeviceSize,
) {
	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  = .STORAGE_BUFFER,
		pBufferInfo     = &buffer_info,
	}

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