#include "static-ui.hpp"

#include "dynamic-array.hpp"
#include "game-settings.hpp"
#include "pk.h"
#include "static-plane.hpp"
#include "vendor-glm-include.hpp"
#include "window.hpp"

#include <cassert>
#include <cstdio>
#include <cstring>
#include <vulkan/vulkan_core.h>

struct pke_ui_box_instance_buffer_item {
	glm::mat4 pos_scale;
	glm::vec2 px_scale;
	float depth;
	float padding[1];
};

struct pke_ui_master {
	pk_membucket *bkt;
	pke_ui_box **root_boxes;
	pke_ui_box_count_T h_root_boxes;
	pke_ui_box_count_T r_root_boxes;
	pke_ui_graphics_bindings bindings;
	bool should_recalc_ui = false;
	glm::vec2 px_scale;
} pke_ui_master;

void pke_ui_init() {
	pke_ui_master.bkt = pk_bucket_create("pke ui", PK_DEFAULT_BUCKET_SIZE, false);
	pke_ui_master.root_boxes = pk_new<pke_ui_box*>(1, pke_ui_master.bkt);
	pke_ui_master.h_root_boxes = 0;
	pke_ui_master.r_root_boxes = 1;
	pke_ui_master.bindings = {};
}

void pke_ui_init_bindings() {
	pkvk_buffer_create_data create_data{};
	create_data.buffer_byte_length[0] = sizeof(pkeIntrinsicsPlane.vert[0]) * 4;
	create_data.buffer_byte_length[1] = sizeof(pkeIntrinsicsPlane.uv[0]) * 4;
	create_data.buffer_byte_length[2] = sizeof(pkeIntrinsicsPlane.index[0]) * 6;
	create_data.buffer_byte_length[3] = sizeof(pke_ui_box_instance_buffer_item);
	create_data.src_data[0] = &pkeIntrinsicsPlane.vert;
	create_data.src_data[1] = &pkeIntrinsicsPlane.uv;
	create_data.src_data[2] = &pkeIntrinsicsPlane.index;
	create_data.src_data[3] = nullptr;
	create_data.n_buffers = 4;
	create_data.index_index = 2;
	create_data.index_instance = 3;

	pkvk_buffer_create_data_out out{};
	pkvk_buffer_create(&create_data, &out);

	pke_ui_master.bindings.bd_vertex.buffer = out.buffers[0];
	pke_ui_master.bindings.bd_vertex.firstBinding = 0;
	pke_ui_master.bindings.bd_vertex.bindingCount = 1;
	pke_ui_master.bindings.bd_uv.buffer = out.buffers[1];
	pke_ui_master.bindings.bd_uv.firstBinding = 1;
	pke_ui_master.bindings.bd_uv.bindingCount = 1;
	pke_ui_master.bindings.bd_index.buffer = out.buffers[2];
	pke_ui_master.bindings.bd_index.firstBinding = 0;
	pke_ui_master.bindings.bd_index.bindingCount = 1;
	pke_ui_master.bindings.index_count = 6;
	pke_ui_master.bindings.bd_instance.buffer = out.buffers[3];
	pke_ui_master.bindings.bd_instance.firstBinding = 2;
	pke_ui_master.bindings.bd_instance.bindingCount = 1;
	pke_ui_master.bindings.deviceMemoryVert = out.device_memory_vertex;
	pke_ui_master.bindings.deviceMemoryInst = out.device_memory_instance;

	pke_ui_master.bindings.instance_buffer_max_count = out.memory_requirements_instance.size / sizeof(pke_ui_box_instance_buffer_item);
}

/* UI layout notes
Some restrictions:
- children cannot change the size of a parent, parents must be sized manually
	- e.g.: exact font width
	- consider writing a method to calculate text
- Children can ONLY be horizontal or vertical
	- this means that tables or similar need to consist of rows of columns
*/

struct pke_ui_flex_params {
	float px_per_unit;
	float unit_total;
};

void pke_ui_calc_px(DynArray<pke_ui_box_instance_buffer_item> &buffer, pke_ui_flex_params *flex_params, pke_ui_box *box) {
	assert(box != nullptr);
	glm::vec2 px_size;
	glm::vec2 px_min_size;
	glm::vec2 px_max_size;
	glm::vec2 parent_size_padded;
	glm::vec2 parent_pos_and_offset;

	assert(box->pos_top_left_x >= 0.0);
	assert(box->pos_top_left_y >= 0.0);
	assert(box->max_width >= 0.0);
	assert(box->max_height >= 0.0);
	assert(box->min_width >= 0.0);
	assert(box->min_height >= 0.0);
	assert(box->min_width <= box->max_width);
	assert(box->min_height <= box->max_height);

	if (box->internal.parent != nullptr) {
		parent_pos_and_offset.x = box->internal.parent->internal.px_corner.x
			+ box->internal.parent->internal.px_padding_l;
		parent_pos_and_offset.y = box->internal.parent->internal.px_corner.y
			+ box->internal.parent->internal.px_padding_t;
		parent_size_padded.x = box->internal.parent->internal.px_size.x;
		parent_size_padded.y = box->internal.parent->internal.px_size.y;
		// built-in padding
		parent_size_padded.x -= box->internal.parent->internal.px_padding_l
			+ box->internal.parent->internal.px_padding_r;
		parent_size_padded.y -= box->internal.parent->internal.px_padding_t
			+ box->internal.parent->internal.px_padding_b;
		px_size = parent_size_padded;
	} else {
		parent_pos_and_offset = glm::vec2(0);
		parent_size_padded = glm::vec2(Extent.width, Extent.height);
		px_size = glm::vec2(Extent.width, Extent.height);
	}

	px_min_size.x = box->min_width;
	px_min_size.y = box->min_height;
	px_max_size.x = box->max_width;
	px_max_size.y = box->max_height;
	box->internal.px_corner.x = 0
		+ parent_pos_and_offset.x
		+ box->pos_top_left_x;
	box->internal.px_corner.y = 0
		+ parent_pos_and_offset.y
		+ box->pos_top_left_y;

	if (PK_HAS_FLAG(box->flags, PKE_UI_BOX_FLAG_POSITION_TYPE_FLEX)) {
		assert(flex_params != nullptr);
		assert(box->internal.parent != nullptr);
		if (box->internal.parent->flex_direction == 0) {
			px_size.x = flex_params->px_per_unit * box->flex_weight;
		} else {
			px_size.y = flex_params->px_per_unit * box->flex_weight;
		}
	} else if (PK_HAS_FLAG(box->flags, PKE_UI_BOX_FLAG_POSITION_TYPE_STATIC)) {
		px_size.x -= box->pos_top_left_x;
		px_size.y -= box->pos_top_left_y;
	} else {
		assert(box->pos_top_left_x < 1.0);
		assert(box->pos_top_left_y < 1.0);
		box->internal.px_corner.x = parent_pos_and_offset.x;
		box->internal.px_corner.y = parent_pos_and_offset.y;
		float px_left = px_size.x * box->pos_top_left_x;
		float px_top = px_size.y * box->pos_top_left_y;
		box->internal.px_corner.x += px_left;
		box->internal.px_corner.y += px_top;
		px_size.x -= px_left;
		px_size.y -= px_top;
		px_min_size.x = box->min_width * parent_size_padded.x;
		px_min_size.y = box->min_height * parent_size_padded.y;
		px_max_size.x = box->max_width * parent_size_padded.x;
		px_max_size.y = box->max_height * parent_size_padded.y;
	}

	px_size = glm::clamp(px_size, px_min_size, px_max_size);

	// built-in padding
	box->internal.px_padding_l = built_in_offset;
	box->internal.px_padding_b = built_in_offset;
	box->internal.px_padding_r = built_in_offset;
	box->internal.px_padding_t = built_in_offset;

	if (PK_HAS_FLAG(box->flags, PKE_UI_BOX_FLAG_CENTER_HORIZONTAL)) {
		if (parent_size_padded.x > px_size.x) {
			box->internal.px_corner.x += (parent_size_padded.x - px_size.x) / 2.0;
		}
	}

	if (PK_HAS_FLAG(box->flags, PKE_UI_BOX_FLAG_CENTER_VERTICAL)) {
		if (parent_size_padded.y > px_size.y) {
			box->internal.px_corner.y += (parent_size_padded.y - px_size.y) / 2.0;
		}
	}

	box->internal.px_size.x = px_size.x;
	box->internal.px_size.y = px_size.y;
	if (box->internal.parent != nullptr) {
		// TODO change this based on the parent's 'direction'
		// 2025-03-13 JCB
		// optional: auto-wrapping boxes possible if we keep track of two sizes:
		// - remaining in active row/column
		// - unused space
		// box->internal.parent->internal.px_padding_l += px_size.x;
		box->internal.parent->internal.px_padding_t += px_size.y;
	}

	if (PK_HAS_FLAG(box->flags, PKE_UI_BOX_FLAG_VISIBILITY_INVISIBLE)) {
		return;
	}

	// update buffer
	glm::vec3 scale = glm::vec3(box->internal.px_size.x / (float)Extent.width, box->internal.px_size.y / (float)Extent.height, 1.0);
	glm::vec3 translate = glm::vec3(1);
	// ((val) - (width/2)) / (width/2)
	// place left line of box at edge of screen
	translate.x = (box->internal.px_size.x / 2.0);
	// box position
	translate.x += box->internal.px_corner.x;
	translate.x -= ((float)Extent.width / 2.0);
	translate.x /= ((float)Extent.width / 2.0);
	// ((val) - (height/2)) / (height/2)
	// place top line of box at edge of screen
	translate.y = (box->internal.px_size.y / 2.0);
	// box position
	translate.y += box->internal.px_corner.y;
	translate.y -= ((float)Extent.height / 2.0);
	translate.y /= ((float)Extent.height / 2.0);
	translate.z = 0;

	pke_ui_box_instance_buffer_item *buffer_item = &buffer.Push();
	buffer_item->pos_scale = glm::translate(glm::mat4(1), translate);
	buffer_item->pos_scale = glm::scale(buffer_item->pos_scale, scale);
	buffer_item->px_scale.x = (2.0 / (float)Extent.width);
	buffer_item->px_scale.y = (2.0 / (float)Extent.height);
	buffer_item->depth = (float)box->layer;
}

void pke_ui_recalc_sizes_recursive(DynArray<pke_ui_box_instance_buffer_item> &arr, pke_ui_box *box, uint8_t depth = 0) {
	uint64_t flags_masked;
	uint8_t flex_count = 0;
	pke_ui_flex_params flex_params{};

	flags_masked = box->flags & (PKE_UI_BOX_FLAG_POSITION_TYPE_ALL);
	if (flags_masked == 0b000 || flags_masked == 0b011 || flags_masked == 0b101 || flags_masked == 0b110 || flags_masked == 0b111) {
		fprintf(stderr, "[%s] ui box invalid flags: position type", __FILE__);
		return;
	}

	for (pke_ui_box_count_T i = 0; i < box->internal.h_children; ++i) {
		if (PK_HAS_FLAG(box->internal.children[i]->flags, PKE_UI_BOX_FLAG_POSITION_TYPE_FLEX)) {
			flex_count += 1;
			flex_params.unit_total += box->internal.children[i]->flex_weight;
		}
	}

	if (flex_count != 0 && flex_count != box->internal.h_children) {
		fprintf(stderr, "[%s] ui box invalid flags: if one child is type FLEX, then all children must be type FLEX", __FILE__);
		return;
	}
	if (flex_count != 0) {
		flex_params.px_per_unit =
			box->flex_direction == 0
			? box->internal.px_size.x / flex_params.unit_total
			: box->internal.px_size.y / flex_params.unit_total;
	}

	for (pke_ui_box_count_T i = 0; i < box->internal.h_children; ++i) {
		pke_ui_calc_px(arr, &flex_params, box->internal.children[i]);
		pke_ui_recalc_sizes_recursive(arr, box->internal.children[i], depth + 1);
	}
}

void pke_ui_update_instance_buffer(DynArray<pke_ui_box_instance_buffer_item> &arr) {
	VkResult vkResult;
	if (arr.Count() > pke_ui_master.bindings.instance_buffer_max_count) {
		VkBuffer newBuffer;
		VkBufferCreateInfo bufferCI;
		bufferCI.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
		bufferCI.pNext = nullptr;
		bufferCI.flags = {};
		bufferCI.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
		bufferCI.queueFamilyIndexCount = 1;
		bufferCI.pQueueFamilyIndices = &graphicsFamilyIndex;
		bufferCI.usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT;
		bufferCI.size = sizeof(pke_ui_box_instance_buffer_item) * arr.Count();
		vkResult = vkCreateBuffer(vkDevice, &bufferCI, vkAllocator, &newBuffer);
		assert(vkResult == VK_SUCCESS);
		VkMemoryRequirements vkMemReqs;
		vkGetBufferMemoryRequirements(vkDevice, newBuffer, &vkMemReqs);
		assert(sizeof(pke_ui_box_instance_buffer_item) % vkMemReqs.alignment == 0);

		vkDestroyBuffer(vkDevice, newBuffer, vkAllocator);
		newBuffer = VK_NULL_HANDLE;

		bufferCI.size = vkMemReqs.size + (vkMemReqs.alignment - (vkMemReqs.size & vkMemReqs.alignment));
		vkResult = vkCreateBuffer(vkDevice, &bufferCI, vkAllocator, &newBuffer);
		assert(vkResult == VK_SUCCESS);

		VkDeviceMemory new_memory;
		VkMemoryAllocateInfo vkMemoryAllocateInfo;
		vkMemoryAllocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
		vkMemoryAllocateInfo.pNext = nullptr;
		vkMemoryAllocateInfo.allocationSize = bufferCI.size;
		vkMemoryAllocateInfo.memoryTypeIndex = FindMemoryTypeIndex(vkMemReqs.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);

		vkResult = vkAllocateMemory(vkDevice, &vkMemoryAllocateInfo, vkAllocator, &new_memory);
		assert(vkResult == VK_SUCCESS);

		vkResult = vkBindBufferMemory(vkDevice, newBuffer, new_memory, 0);
		assert(vkResult == VK_SUCCESS);
		if (pke_ui_master.bindings.bd_instance.buffer != VK_NULL_HANDLE) {
			vkDestroyBuffer(vkDevice, pke_ui_master.bindings.bd_instance.buffer, vkAllocator);
		}
		if (pke_ui_master.bindings.deviceMemoryInst != VK_NULL_HANDLE) {
			vkFreeMemory(vkDevice, pke_ui_master.bindings.deviceMemoryInst, vkAllocator);
		}
		pke_ui_master.bindings.bd_instance.buffer = newBuffer;
		pke_ui_master.bindings.deviceMemoryInst = new_memory;
		pke_ui_master.bindings.instance_buffer_max_count = arr.Count();
	}
	PKVK_TmpBufferDetails tmpBufferDetails{};
	PKVK_BeginBuffer(graphicsFamilyIndex, sizeof(pke_ui_box_instance_buffer_item) * arr.Count(), tmpBufferDetails);
	{
		VkCommandBufferBeginInfo vkCommandBufferBeginInfo;
		vkCommandBufferBeginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
		vkCommandBufferBeginInfo.pNext = nullptr;
		// TODO consider single-use?
		vkCommandBufferBeginInfo.flags = 0;
		vkCommandBufferBeginInfo.pInheritanceInfo = nullptr;
		vkResult = vkBeginCommandBuffer(tmpBufferDetails.cmdBuffer, &vkCommandBufferBeginInfo);
		assert(vkResult == VK_SUCCESS);

		memcpy(tmpBufferDetails.deviceData, arr.GetPtr(), sizeof(pke_ui_box_instance_buffer_item) * arr.Count());

		VkBufferCopy vk_buffer_copy{};
		vk_buffer_copy.srcOffset = 0;
		vk_buffer_copy.dstOffset = 0;
		vk_buffer_copy.size = sizeof(pke_ui_box_instance_buffer_item) * arr.Count();
		vkCmdCopyBuffer(tmpBufferDetails.cmdBuffer, tmpBufferDetails.buffer, pke_ui_master.bindings.bd_instance.buffer, 1, &vk_buffer_copy);

		vkResult = vkEndCommandBuffer(tmpBufferDetails.cmdBuffer);
		assert(vkResult == VK_SUCCESS);

		VkSubmitInfo submitInfo;
		submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
		submitInfo.pNext = nullptr;
		submitInfo.waitSemaphoreCount = 0;
		submitInfo.pWaitSemaphores = nullptr;
		submitInfo.pWaitDstStageMask = nullptr;
		submitInfo.commandBufferCount = 1;
		submitInfo.pCommandBuffers = &tmpBufferDetails.cmdBuffer;
		submitInfo.signalSemaphoreCount = 0;
		submitInfo.pSignalSemaphores = nullptr;
		vkResult = vkQueueSubmit(tmpBufferDetails.queue, 1, &submitInfo, nullptr);
		assert(vkResult == VK_SUCCESS);
		vkResult = vkQueueWaitIdle(tmpBufferDetails.queue);
		assert(vkResult == VK_SUCCESS);
	}
	PKVK_EndBuffer(tmpBufferDetails);
	pke_ui_master.bindings.instance_counter = arr.Count();
}

void pke_ui_tick(double delta) {
	(void)delta;
	if (pke_ui_master.h_root_boxes == 0) return;
	if (pke_ui_master.should_recalc_ui == true || pkeSettings.rt.was_framebuffer_resized == true) {
		DynArray<pke_ui_box_instance_buffer_item> arr;
		pke_ui_master.should_recalc_ui = false;
		pke_ui_master.px_scale = glm::vec2(
			2.0 / (float)Extent.width,
			2.0 / (float)Extent.height
		);
		for (pke_ui_box_count_T i = 0; i < pke_ui_master.h_root_boxes; ++i) {
			pke_ui_box *box = pke_ui_master.root_boxes[i];
			pke_ui_calc_px(arr, nullptr, box);
			pke_ui_recalc_sizes_recursive(arr, box, 0);
		}
		pke_ui_update_instance_buffer(arr);
	}
}

void pke_ui_teardown_box_recursive(pke_ui_box *box) {
	for (pke_ui_box_count_T i = 0; i < box->internal.h_children; ++i) {
		pke_ui_teardown_box_recursive(box->internal.children[i]);
	}
	if (box->internal.children != nullptr) {
		pk_delete<pke_ui_box *>(box->internal.children, box->internal.r_children);
	}
}

void pke_ui_teardown() {
	for (pke_ui_box_count_T i = 0; i < pke_ui_master.h_root_boxes; ++i) {
		pke_ui_teardown_box_recursive(pke_ui_master.root_boxes[i]);
	}
	pk_delete<pke_ui_box *>(pke_ui_master.root_boxes, pke_ui_master.r_root_boxes);
	pk_bucket_destroy(pke_ui_master.bkt);
	pke_ui_master.bkt = nullptr;
	pke_ui_master.root_boxes = nullptr;
	pke_ui_master.h_root_boxes = 0;
	pke_ui_master.r_root_boxes = 0;

	if (pke_ui_master.bindings.bd_vertex.buffer != VK_NULL_HANDLE)
		vkDestroyBuffer(vkDevice, pke_ui_master.bindings.bd_vertex.buffer, vkAllocator);
	pke_ui_master.bindings.bd_vertex.buffer = VK_NULL_HANDLE;
	pke_ui_master.bindings.bd_vertex.firstBinding = 0;
	pke_ui_master.bindings.bd_vertex.bindingCount = 0;
	pke_ui_master.bindings.bd_vertex.offsets[0] = 0;

	if (pke_ui_master.bindings.bd_uv.buffer != VK_NULL_HANDLE)
		vkDestroyBuffer(vkDevice, pke_ui_master.bindings.bd_uv.buffer, vkAllocator);
	pke_ui_master.bindings.bd_uv.buffer = VK_NULL_HANDLE;
	pke_ui_master.bindings.bd_uv.firstBinding = 0;
	pke_ui_master.bindings.bd_uv.bindingCount = 0;
	pke_ui_master.bindings.bd_uv.offsets[0] = 0;

	if (pke_ui_master.bindings.bd_index.buffer != VK_NULL_HANDLE)
		vkDestroyBuffer(vkDevice, pke_ui_master.bindings.bd_index.buffer, vkAllocator);
	pke_ui_master.bindings.bd_index.buffer = VK_NULL_HANDLE;
	pke_ui_master.bindings.bd_index.bindingCount = 0;
	pke_ui_master.bindings.bd_index.offsets[0] = 0;
	pke_ui_master.bindings.index_count = 0;

	if (pke_ui_master.bindings.bd_instance.buffer != VK_NULL_HANDLE)
		vkDestroyBuffer(vkDevice, pke_ui_master.bindings.bd_instance.buffer, vkAllocator);
	pke_ui_master.bindings.bd_instance.buffer = VK_NULL_HANDLE;
	pke_ui_master.bindings.bd_instance.firstBinding = 0;
	pke_ui_master.bindings.bd_instance.bindingCount = 0;
	pke_ui_master.bindings.instance_counter = 0;
	pke_ui_master.bindings.instance_buffer_max_count = 0;
	pke_ui_master.bindings.bd_instance.offsets[0] = 0;

	if (pke_ui_master.bindings.deviceMemoryInst != VK_NULL_HANDLE)
		vkFreeMemory(vkDevice, pke_ui_master.bindings.deviceMemoryInst, vkAllocator);
	pke_ui_master.bindings.deviceMemoryInst = VK_NULL_HANDLE;

	if (pke_ui_master.bindings.deviceMemoryVert != VK_NULL_HANDLE)
		vkFreeMemory(vkDevice, pke_ui_master.bindings.deviceMemoryVert, vkAllocator);
	pke_ui_master.bindings.deviceMemoryVert = VK_NULL_HANDLE;
}

pke_ui_box *pke_ui_box_new_root() {
	if (pke_ui_master.h_root_boxes == pke_ui_master.r_root_boxes) {
		pke_ui_box_count_T prev_r_root_boxes = pke_ui_master.r_root_boxes;
		pke_ui_master.r_root_boxes *= 1.5;
		pke_ui_box **boxes = pk_new<pke_ui_box*>(pke_ui_master.r_root_boxes);
		for (pke_ui_box_count_T i = 0; i < pke_ui_master.h_root_boxes; ++i) {
			boxes[i] = pke_ui_master.root_boxes[i];
		}
		if (pke_ui_master.root_boxes != nullptr) pk_delete<pke_ui_box*>(pke_ui_master.root_boxes, prev_r_root_boxes);
		pke_ui_master.root_boxes = boxes;
	}
	pke_ui_box *box = pk_new<pke_ui_box>(pke_ui_master.bkt);
	memset(box, 0, sizeof(pke_ui_box));
	pke_ui_master.root_boxes[pke_ui_master.h_root_boxes] = box;
	pke_ui_master.h_root_boxes += 1;
	pke_ui_master.should_recalc_ui = true;
	return box;
}

pke_ui_box *pke_ui_box_new_child(pke_ui_box *parent) {
	assert(parent != nullptr);
	if (parent->internal.h_children == parent->internal.r_children) {
		pke_ui_box_count_T prev_r_children = parent->internal.r_children;
		parent->internal.r_children = PK_MAX(parent->internal.r_children * 1.5, 2);
		pke_ui_box **boxes = pk_new<pke_ui_box*>(parent->internal.r_children);
		for (pke_ui_box_count_T i = 0; i < parent->internal.h_children; ++i) {
			boxes[i] = parent->internal.children[i];
		}
		if (parent->internal.children != nullptr) {
			pk_delete<pke_ui_box*>(parent->internal.children, prev_r_children);
		}
		parent->internal.children = boxes;
	}
	pke_ui_box *box = pk_new<pke_ui_box>(pke_ui_master.bkt);
	memset(box, 0, sizeof(pke_ui_box));
	parent->internal.children[parent->internal.h_children] = box;
	parent->internal.h_children += 1;
	box->internal.parent = parent;
	pke_ui_master.should_recalc_ui = true;
	return box;
}

pke_ui_graphics_bindings *pke_ui_get_graphics_bindings() {
	return &pke_ui_master.bindings;
}