#include "vulkan/basalt_device.h"

basalt::DeviceDesigner::DeviceDesigner(basalt::Context& ctx, basalt::PhysicalDevice& phy) :
	ctx(ctx), phys(phy), req_features(VkPhysicalDeviceFeatures{}),
	req_queues(MEMORY_TAG_CLASS_DYNARRAY | MEMORY_TAG_ZONE_ENGINE | MEMORY_TAG_ALIGN_ANY),
	req_extensions(phy.enabled_extensions),
	req_layers(MEMORY_TAG_CLASS_DYNARRAY | MEMORY_TAG_ZONE_ENGINE | MEMORY_TAG_ALIGN_ANY),
	queue_priorities(MEMORY_TAG_CLASS_DYNARRAY | MEMORY_TAG_ZONE_ENGINE | MEMORY_TAG_ALIGN_ANY)
{}

basalt::DeviceDesigner& basalt::DeviceDesigner::request_features(VkPhysicalDeviceFeatures features)
{
	VkBool32* cur = (VkBool32*)(&this->req_features);
	const VkBool32* req = (VkBool32*)(&features);
	for (u16 i = 0; i < sizeof(VkPhysicalDeviceFeatures) / sizeof(VkBool32); ++i)
		cur[i] |= req[i];
	return *this;
}

basalt::DeviceDesigner& basalt::DeviceDesigner::request_extensions(const basalt::darray<const char*>& extension)
{
	this->req_extensions.push_back(extension);
	return *this;
}

basalt::DeviceDesigner& basalt::DeviceDesigner::request_extensions(const char** extensions, const u32 num_extensions)
{
	this->req_extensions.push_back(extensions, num_extensions);
	return *this;
}

basalt::DeviceDesigner& basalt::DeviceDesigner::request_extension(const char* extension)
{
	this->req_extensions.push_back(extension);
	return *this;
}

basalt::DeviceDesigner& basalt::DeviceDesigner::request_layers(const basalt::darray<const char*>&layers)
{
	this->req_layers.push_back(layers);
	return *this;
}

basalt::DeviceDesigner& basalt::DeviceDesigner::request_layers(const char** layers, const u32 num_layers)
{
	this->req_layers.push_back(layers, num_layers);
	return *this;
}

basalt::DeviceDesigner& basalt::DeviceDesigner::request_layer(const char* layer)
{
	this->req_layers.push_back(layer);
	return *this;
}

basalt::DeviceDesigner& basalt::DeviceDesigner::request_queue(u32 family_index, const float priority, VkDeviceQueueCreateFlags flags)
{
	const u32 priority_ref = this->queue_priorities.m_nelements;
	this->queue_priorities.push_back(priority);
	VkDeviceQueueCreateInfo ci = { VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO };
	ci.flags = flags;
	ci.pNext = VK_NULL_HANDLE;
	ci.pQueuePriorities = (float*)priority_ref;
	ci.queueCount = 1;
	ci.queueFamilyIndex = family_index;
	this->req_queues.push_back(ci);
	return *this;
}

basalt::DeviceDesigner& basalt::DeviceDesigner::request_queues(u32 family_index, const basalt::darray<float>& queue_priorities, VkDeviceQueueCreateFlags flags)
{
	const u32 priority_ref = this->queue_priorities.m_nelements;
	this->queue_priorities.push_back(queue_priorities);
	VkDeviceQueueCreateInfo ci = { VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO };
	ci.flags = flags;
	ci.pNext = VK_NULL_HANDLE;
	ci.pQueuePriorities = (float*)priority_ref;
	ci.queueCount = queue_priorities.size();
	ci.queueFamilyIndex = family_index;
	this->req_queues.push_back(ci);
	return *this;
}

basalt::Device basalt::DeviceDesigner::design(void)
{
	VkDevice logical = VK_NULL_HANDLE;

	// Resolve the relative references for queue priorities in the queue create infos
	for (VkDeviceQueueCreateInfo& queue_ci : this->req_queues)
		queue_ci.pQueuePriorities = ((u32)queue_ci.pQueuePriorities) + this->queue_priorities.begin();

	VkDeviceCreateInfo ci = {VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO};
	ci.flags = 0;
	ci.queueCreateInfoCount = this->req_queues.m_nelements;
	ci.pQueueCreateInfos = this->req_queues.m_pdata;
	ci.pEnabledFeatures = &this->req_features;
	ci.enabledExtensionCount = this->req_extensions.m_nelements;
	ci.ppEnabledExtensionNames = this->req_extensions.m_pdata;
	ci.enabledLayerCount = this->req_layers.m_nelements;
	ci.ppEnabledLayerNames = this->req_layers.m_pdata;

	VkResult err = VK_SUCCESS;
	VK_ASSERT(vkCreateDevice(this->phys, &ci, this->ctx.vk_alloc, &logical), "Failed to create logical device\n\tAt %s:%d\n\tError %s\n");
	basalt::Device dev(this->ctx, this->phys, logical);
	
	u32 num_total_queues = 0;
	for (u32 i = 0; i < this->req_queues.m_nelements; ++i)
		num_total_queues += this->req_queues[i].queueCount;
	dev.queues.reserve(num_total_queues);

	u32 running_total = 0;
	for (u32 i = 0; i < this->req_queues.m_nelements; ++i)
		for (u32 j = 0; j < this->req_queues[i].queueCount; ++j)
			vkGetDeviceQueue(logical, this->req_queues[i].queueFamilyIndex, j, &dev.queues[running_total++].queue);
	dev.queues.m_nelements = num_total_queues;

	return std::move(dev);
}

basalt::Device::Device(const Device& src) :
	ctx(src.ctx), logical(src.logical), phys(src.phys), queues(src.queues), should_free(false)
{}

basalt::Device& basalt::Device::operator=(const Device & src)
{
	if (&src == this)
		return *this;
	if (src.logical != this->logical && this->logical != VK_NULL_HANDLE && this->should_free)
		this->~Device();
	this->ctx = src.ctx;
	this->logical = src.logical;
	this->phys = src.phys;
	this->queues = darray<basalt::Queue>(src.queues);
	this->should_free = false;
	return *this;
}

basalt::Device::Device(Device&& other) noexcept :
	ctx(other.ctx), logical(other.logical), phys(other.phys), queues(std::move(other.queues)), should_free(other.should_free)
{
	other.ctx = nullptr;
	other.logical = VK_NULL_HANDLE;
	other.phys = nullptr;
	other.should_free = false;
}

basalt::Device& basalt::Device::operator=(Device&& other) noexcept
{
	if (&other == this)
		return *this;
	if (this->logical != other.logical && this->logical != VK_NULL_HANDLE && this->should_free)
		this->~Device();
	
	this->ctx = other.ctx;
	this->logical = other.logical;
	this->phys = other.phys;
	this->should_free = other.should_free;
	this->queues = std::move(other.queues);

	other.ctx = nullptr;
	other.logical = VK_NULL_HANDLE;
	other.phys = nullptr;
	other.should_free = false;

	return *this;
}

basalt::Device::Device(basalt::Context& ctx, basalt::PhysicalDevice& physical, VkDevice logical) :
	ctx(&ctx), phys(&physical), logical(logical), queues(MEMORY_TAG_CLASS_ARRAY | MEMORY_TAG_ZONE_ENGINE | MEMORY_TAG_ALIGN_ANY), should_free(true)
{}

basalt::Device::~Device()
{
	if (this->should_free && this->logical != VK_NULL_HANDLE)
		vkDestroyDevice(this->logical, this->ctx->vk_alloc);
}

basalt::Device::operator VkPhysicalDevice()
{ return this->phys->phys; }

basalt::Device::operator VkDevice()
{ return this->logical; }