package gpu_core

import "core:log"
import vk "vendor:vulkan"

Gpu_Image :: struct {
	image:  vk.Image,
	memory: vk.DeviceMemory,
	view:   vk.ImageView,
	width:  u32,
	height: u32,
	format: vk.Format,
}

Gpu_Buffer :: struct {
	buffer: vk.Buffer,
	memory: vk.DeviceMemory,
	size:   vk.DeviceSize,
	mapped: rawptr,
}

create_buffer :: proc(
	dev: ^Gpu_Device,
	size: vk.DeviceSize,
	usage: vk.BufferUsageFlags,
	properties: vk.MemoryPropertyFlags,
) -> (buf: Gpu_Buffer, ok: bool) {
	buf_info := vk.BufferCreateInfo{
		sType       = .BUFFER_CREATE_INFO,
		size        = size,
		usage       = usage,
		sharingMode = .EXCLUSIVE,
	}

	result := vk.CreateBuffer(dev.device, &buf_info, nil, &buf.buffer)
	if result != .SUCCESS {
		log.errorf("gpu/core: vkCreateBuffer failed: %v", result)
		return {}, false
	}

	mem_reqs: vk.MemoryRequirements
	vk.GetBufferMemoryRequirements(dev.device, buf.buffer, &mem_reqs)

	type_index, type_ok := find_memory_type(dev, mem_reqs.memoryTypeBits, properties)
	if !type_ok {
		vk.DestroyBuffer(dev.device, buf.buffer, nil)
		return {}, false
	}

	alloc_info := vk.MemoryAllocateInfo{
		sType           = .MEMORY_ALLOCATE_INFO,
		allocationSize  = mem_reqs.size,
		memoryTypeIndex = type_index,
	}

	result = vk.AllocateMemory(dev.device, &alloc_info, nil, &buf.memory)
	if result != .SUCCESS {
		log.errorf("gpu/core: vkAllocateMemory failed: %v", result)
		vk.DestroyBuffer(dev.device, buf.buffer, nil)
		return {}, false
	}

	vk.BindBufferMemory(dev.device, buf.buffer, buf.memory, 0)
	buf.size = size

	return buf, true
}

// Create a device-local buffer by staging through a host-visible buffer
create_buffer_staged :: proc(
	dev: ^Gpu_Device,
	cmd_pool: vk.CommandPool,
	queue: vk.Queue,
	data: rawptr,
	size: int,
	usage: vk.BufferUsageFlags,
) -> (buf: Gpu_Buffer, ok: bool) {
	vk_size := vk.DeviceSize(size)

	// Create staging buffer
	staging, staging_ok := create_buffer(
		dev, vk_size,
		{.TRANSFER_SRC},
		{.HOST_VISIBLE, .HOST_COHERENT},
	)
	if !staging_ok {
		return {}, false
	}
	defer destroy_buffer(dev, &staging)

	// Map and copy data
	mapped: rawptr
	vk.MapMemory(dev.device, staging.memory, 0, vk_size, {}, &mapped)
	copy_raw(mapped, data, size)
	vk.UnmapMemory(dev.device, staging.memory)

	// Create device-local buffer
	result_buf, result_ok := create_buffer(
		dev, vk_size,
		usage + {.TRANSFER_DST},
		{.DEVICE_LOCAL},
	)
	if !result_ok {
		return {}, false
	}
	buf = result_buf

	// Copy via command buffer
	cmd, cmd_ok := begin_single_time_commands(dev, cmd_pool)
	if !cmd_ok {
		destroy_buffer(dev, &buf)
		return {}, false
	}

	region := vk.BufferCopy{
		size = vk_size,
	}
	vk.CmdCopyBuffer(cmd, staging.buffer, buf.buffer, 1, &region)

	end_single_time_commands(dev, cmd_pool, queue, cmd)

	return buf, true
}

destroy_buffer :: proc(dev: ^Gpu_Device, buf: ^Gpu_Buffer) {
	if buf.buffer != 0 {
		vk.DestroyBuffer(dev.device, buf.buffer, nil)
		buf.buffer = 0
	}
	if buf.memory != 0 {
		vk.FreeMemory(dev.device, buf.memory, nil)
		buf.memory = 0
	}
}

find_memory_type :: proc(dev: ^Gpu_Device, type_filter: u32, properties: vk.MemoryPropertyFlags) -> (u32, bool) {
	mem_props: vk.PhysicalDeviceMemoryProperties
	vk.GetPhysicalDeviceMemoryProperties(dev.physical_device, &mem_props)

	for i in 0..<mem_props.memoryTypeCount {
		if (type_filter & (1 << i)) != 0 {
			if properties <= mem_props.memoryTypes[i].propertyFlags {
				return i, true
			}
		}
	}

	log.error("gpu/core: failed to find suitable memory type")
	return 0, false
}

// --- Image operations ---

create_image :: proc(
	dev: ^Gpu_Device,
	width, height: u32,
	format: vk.Format,
	tiling: vk.ImageTiling,
	usage: vk.ImageUsageFlags,
	properties: vk.MemoryPropertyFlags,
) -> (img: Gpu_Image, ok: bool) {
	image_info := vk.ImageCreateInfo{
		sType       = .IMAGE_CREATE_INFO,
		imageType   = .D2,
		extent      = {width, height, 1},
		mipLevels   = 1,
		arrayLayers = 1,
		format      = format,
		tiling      = tiling,
		initialLayout = .UNDEFINED,
		usage       = usage,
		sharingMode = .EXCLUSIVE,
		samples     = {._1},
	}

	result := vk.CreateImage(dev.device, &image_info, nil, &img.image)
	if result != .SUCCESS {
		log.errorf("gpu/core: vkCreateImage failed: %v", result)
		return {}, false
	}

	mem_reqs: vk.MemoryRequirements
	vk.GetImageMemoryRequirements(dev.device, img.image, &mem_reqs)

	type_index, type_ok := find_memory_type(dev, mem_reqs.memoryTypeBits, properties)
	if !type_ok {
		vk.DestroyImage(dev.device, img.image, nil)
		return {}, false
	}

	alloc_info := vk.MemoryAllocateInfo{
		sType           = .MEMORY_ALLOCATE_INFO,
		allocationSize  = mem_reqs.size,
		memoryTypeIndex = type_index,
	}

	result = vk.AllocateMemory(dev.device, &alloc_info, nil, &img.memory)
	if result != .SUCCESS {
		log.errorf("gpu/core: vkAllocateMemory for image failed: %v", result)
		vk.DestroyImage(dev.device, img.image, nil)
		return {}, false
	}

	vk.BindImageMemory(dev.device, img.image, img.memory, 0)
	img.width = width
	img.height = height
	img.format = format

	return img, true
}

create_image_view :: proc(dev: ^Gpu_Device, image: vk.Image, format: vk.Format, aspect: vk.ImageAspectFlags = {.COLOR}) -> (view: vk.ImageView, ok: bool) {
	view_info := vk.ImageViewCreateInfo{
		sType    = .IMAGE_VIEW_CREATE_INFO,
		image    = image,
		viewType = .D2,
		format   = format,
		subresourceRange = {
			aspectMask     = aspect,
			baseMipLevel   = 0,
			levelCount     = 1,
			baseArrayLayer = 0,
			layerCount     = 1,
		},
	}

	result := vk.CreateImageView(dev.device, &view_info, nil, &view)
	if result != .SUCCESS {
		log.errorf("gpu/core: vkCreateImageView failed: %v", result)
		return {}, false
	}

	return view, true
}

transition_image_layout :: proc(
	dev: ^Gpu_Device,
	cmd_pool: vk.CommandPool,
	queue: vk.Queue,
	image: vk.Image,
	old_layout: vk.ImageLayout,
	new_layout: vk.ImageLayout,
) -> bool {
	cmd, cmd_ok := begin_single_time_commands(dev, cmd_pool)
	if !cmd_ok {
		return false
	}

	barrier := vk.ImageMemoryBarrier{
		sType            = .IMAGE_MEMORY_BARRIER,
		oldLayout        = old_layout,
		newLayout        = new_layout,
		srcQueueFamilyIndex = vk.QUEUE_FAMILY_IGNORED,
		dstQueueFamilyIndex = vk.QUEUE_FAMILY_IGNORED,
		image            = image,
		subresourceRange = {
			aspectMask     = {.COLOR},
			baseMipLevel   = 0,
			levelCount     = 1,
			baseArrayLayer = 0,
			layerCount     = 1,
		},
	}

	src_stage, dst_stage: vk.PipelineStageFlags

	if old_layout == .UNDEFINED && new_layout == .TRANSFER_DST_OPTIMAL {
		barrier.srcAccessMask = {}
		barrier.dstAccessMask = {.TRANSFER_WRITE}
		src_stage = {.TOP_OF_PIPE}
		dst_stage = {.TRANSFER}
	} else if old_layout == .TRANSFER_DST_OPTIMAL && new_layout == .SHADER_READ_ONLY_OPTIMAL {
		barrier.srcAccessMask = {.TRANSFER_WRITE}
		barrier.dstAccessMask = {.SHADER_READ}
		src_stage = {.TRANSFER}
		dst_stage = {.FRAGMENT_SHADER}
	} else {
		log.error("gpu/core: unsupported image layout transition")
		return false
	}

	vk.CmdPipelineBarrier(cmd, src_stage, dst_stage, {}, 0, nil, 0, nil, 1, &barrier)

	end_single_time_commands(dev, cmd_pool, queue, cmd)
	return true
}

copy_buffer_to_image :: proc(
	dev: ^Gpu_Device,
	cmd_pool: vk.CommandPool,
	queue: vk.Queue,
	buffer: vk.Buffer,
	image: vk.Image,
	width, height: u32,
) -> bool {
	cmd, cmd_ok := begin_single_time_commands(dev, cmd_pool)
	if !cmd_ok {
		return false
	}

	region := vk.BufferImageCopy{
		bufferOffset      = 0,
		bufferRowLength   = 0,
		bufferImageHeight = 0,
		imageSubresource  = {
			aspectMask     = {.COLOR},
			mipLevel       = 0,
			baseArrayLayer = 0,
			layerCount     = 1,
		},
		imageOffset = {0, 0, 0},
		imageExtent = {width, height, 1},
	}

	vk.CmdCopyBufferToImage(cmd, buffer, image, .TRANSFER_DST_OPTIMAL, 1, &region)

	end_single_time_commands(dev, cmd_pool, queue, cmd)
	return true
}

create_sampler :: proc(dev: ^Gpu_Device) -> (sampler: vk.Sampler, ok: bool) {
	max_aniso := dev.properties.limits.maxSamplerAnisotropy

	sampler_info := vk.SamplerCreateInfo{
		sType        = .SAMPLER_CREATE_INFO,
		magFilter    = .LINEAR,
		minFilter    = .LINEAR,
		addressModeU = .REPEAT,
		addressModeV = .REPEAT,
		addressModeW = .REPEAT,
		anisotropyEnable = true,
		maxAnisotropy    = max_aniso,
		borderColor      = .INT_OPAQUE_BLACK,
		unnormalizedCoordinates = false,
		compareEnable    = false,
		compareOp        = .ALWAYS,
		mipmapMode       = .LINEAR,
		mipLodBias       = 0,
		minLod           = 0,
		maxLod           = 0,
	}

	result := vk.CreateSampler(dev.device, &sampler_info, nil, &sampler)
	if result != .SUCCESS {
		log.errorf("gpu/core: vkCreateSampler failed: %v", result)
		return {}, false
	}

	return sampler, true
}

destroy_image :: proc(dev: ^Gpu_Device, img: ^Gpu_Image) {
	if img.view != 0 {
		vk.DestroyImageView(dev.device, img.view, nil)
		img.view = 0
	}
	if img.image != 0 {
		vk.DestroyImage(dev.device, img.image, nil)
		img.image = 0
	}
	if img.memory != 0 {
		vk.FreeMemory(dev.device, img.memory, nil)
		img.memory = 0
	}
}

@(private)
copy_raw :: proc(dst, src: rawptr, size: int) {
	d := ([^]u8)(dst)
	s := ([^]u8)(src)
	for i in 0..<size {
		d[i] = s[i]
	}
}