#include "renderer.h"
#include "buffer.h"

#include <string.h>

static inline ssize_t find_memory_type(struct vlkn_renderer *ren, uint32_t filter, VkMemoryPropertyFlags props) {
	VkPhysicalDeviceMemoryProperties mem_props;
	vkGetPhysicalDeviceMemoryProperties(ren->phy_gpus.chosen->gpu, &mem_props);

	for (size_t i = 0; i < mem_props.memoryTypeCount; i++)
		if (filter & (1 << i) && (mem_props.memoryTypes[i].propertyFlags & props) == props)
			return i;

	return -1;
}

VkResult buffer_create(struct vlkn_renderer *ren, VkDeviceSize size,
		VkBufferUsageFlags usage, VkMemoryPropertyFlags props, struct vlkn_buffer *buffer) {
	VkBufferCreateInfo buffer_info = {
		.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
		.size = size,
		.usage = usage,
		.sharingMode = VK_SHARING_MODE_EXCLUSIVE,
	};

	VkResult res = vkCreateBuffer(ren->gpu.device, &buffer_info, NULL, &buffer->buffer);
	if (res != VK_SUCCESS)
		return res;

	VkMemoryRequirements mem_reqs;
	vkGetBufferMemoryRequirements(ren->gpu.device, buffer->buffer, &mem_reqs);

	ssize_t mem_type_index = find_memory_type(ren, mem_reqs.memoryTypeBits, props);
	if (mem_type_index == -1)
		return VK_ERROR_UNKNOWN;

	VkMemoryAllocateInfo alloc_info = {
		.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
		.allocationSize = mem_reqs.size,
		.memoryTypeIndex = mem_type_index
	};

	res = vkAllocateMemory(ren->gpu.device, &alloc_info, NULL, &buffer->memory);
	if (res != VK_SUCCESS)
		return res;

	buffer->size = size;
	return vkBindBufferMemory(ren->gpu.device, buffer->buffer, buffer->memory, 0);
}

VkResult buffer_upload(struct vlkn_renderer *ren, struct vlkn_buffer *buffer,
		size_t offset, size_t size, uint8_t data[size]) {
	struct vlkn_buffer tmp;
	size_t end_size = offset + size;
	// TODO: reallocate buffer?
	if (end_size > buffer->size) {
		end_size = buffer->size - offset;
	}
	VkResult res = buffer_create(ren, end_size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
			VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, &tmp);
	if (res != VK_SUCCESS)
		return res;

	void *map;
	res = vkMapMemory(ren->gpu.device, tmp.memory, 0, end_size, 0, &map);
	if (res != VK_SUCCESS)
		return res;
	memcpy(map, data, size < buffer->size ? size : end_size);
	vkUnmapMemory(ren->gpu.device, tmp.memory);

	VkCommandBuffer cmd = begin_single_command(ren);
	vkCmdCopyBuffer(cmd, tmp.buffer, buffer->buffer, 1, &(VkBufferCopy) { .size = end_size });
	end_single_command(ren, cmd);

	buffer_destroy(ren, &tmp);

	return VK_SUCCESS;
}

void buffer_destroy(struct vlkn_renderer *ren, struct vlkn_buffer *buffer) {
	vkDestroyBuffer(ren->gpu.device, buffer->buffer, NULL);
	vkFreeMemory(ren->gpu.device, buffer->memory, NULL);
}

VkResult image_view_create(struct vlkn_renderer *ren, VkFormat format,
		VkImageAspectFlags aspect, VkImage image, VkImageView *view) {
	VkImageViewCreateInfo view_info = {
		.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
		.image = image,
		.viewType = VK_IMAGE_VIEW_TYPE_2D,
		.format = format,
		.components = {
			.r = VK_COMPONENT_SWIZZLE_IDENTITY,
			.g = VK_COMPONENT_SWIZZLE_IDENTITY,
			.b = VK_COMPONENT_SWIZZLE_IDENTITY,
			.a = VK_COMPONENT_SWIZZLE_IDENTITY,
		},
		.subresourceRange = {
			.aspectMask = aspect,
			.baseMipLevel = 0,
			.levelCount = 1,
			.baseArrayLayer = 0,
			.layerCount = 1
		}
	};

	return vkCreateImageView(ren->gpu.device, &view_info, NULL, view);
}

VkResult image_create(struct vlkn_renderer *ren, struct image_opts opts, struct vlkn_images *image) {
	VkImageCreateInfo image_info = {
		.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
		.imageType = VK_IMAGE_TYPE_2D,
		.extent = {
			.width = opts.extent.width,
			.height = opts.extent.height,
			.depth = 1
		},
		.mipLevels = opts.mip_level,
		.arrayLayers = 1,
		.format = opts.format,
		.tiling = opts.tiling,
		.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
		.usage = opts.usage,
		.sharingMode = VK_SHARING_MODE_EXCLUSIVE,
		.samples = VK_SAMPLE_COUNT_1_BIT
	};

	VkResult res = vkCreateImage(ren->gpu.device, &image_info, NULL, &image->image);
	if (res != VK_SUCCESS)
		return res;

	VkMemoryRequirements mem_reqs;
	vkGetImageMemoryRequirements(ren->gpu.device, image->image, &mem_reqs);

	ssize_t mem_type_index = find_memory_type(ren, mem_reqs.memoryTypeBits, opts.mem_props);
	if (mem_type_index == -1)
		return VK_ERROR_UNKNOWN;

	VkMemoryAllocateInfo alloc_info = {
		.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
		.allocationSize = mem_reqs.size,
		.memoryTypeIndex = mem_type_index
	};

	res = vkAllocateMemory(ren->gpu.device, &alloc_info, NULL, &image->memory);
	if (res != VK_SUCCESS)
		return res;

	res = vkBindImageMemory(ren->gpu.device, image->image, image->memory, 0);
	if (res != VK_SUCCESS)
		return res;

	return image_view_create(ren, opts.format, opts.aspect, image->image, &image->view);
}

void image_destroy(struct vlkn_renderer *ren, struct vlkn_images *image) {
	vkDestroyImageView(ren->gpu.device, image->view, NULL);
	vkDestroyImage(ren->gpu.device, image->image, NULL);
	vkFreeMemory(ren->gpu.device, image->memory, NULL);
}

struct vlkn_mesh vlkn_mesh_upload(struct vlkn_renderer *renderer, size_t vertex_count, struct vlkn_vertex vertices[vertex_count],
		size_t index_count, uint32_t indices[index_count], struct vec3 position) {
	(void) position;
	const VkDeviceSize vertex_size = sizeof(*vertices) * vertex_count;
	const VkDeviceSize index_size = sizeof(*indices) * index_count;

	struct vlkn_mesh mesh;
	buffer_create(renderer, vertex_size, VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
			VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, &mesh.vertex);
	buffer_create(renderer, index_size, VK_BUFFER_USAGE_INDEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
			VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, &mesh.index);

	struct vlkn_buffer tmp;
	buffer_create(renderer, vertex_size + index_size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
			VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, &tmp);

	void *data;
	vkMapMemory(renderer->gpu.device, tmp.memory, 0, vertex_size + index_size, 0, &data);
	memcpy(data, vertices, vertex_size);
	memcpy((char *)data + vertex_size, indices, index_size);
	vkUnmapMemory(renderer->gpu.device, tmp.memory);

	VkCommandBuffer cmd = begin_single_command(renderer);
	vkCmdCopyBuffer(cmd, tmp.buffer, mesh.vertex.buffer, 1, &(VkBufferCopy) { .size = vertex_size });
	vkCmdCopyBuffer(cmd, tmp.buffer, mesh.index.buffer, 1, &(VkBufferCopy) { .size = index_size, .srcOffset = vertex_size });
	end_single_command(renderer, cmd);

	buffer_destroy(renderer, &tmp);
	mesh.index_count = index_count;
	mesh.position = position;

	return mesh;
}