#include "stdafx.h"

#include "VKFramebuffer.h"
#include "vkutils/image.h"
#include "vkutils/image_helpers.h"

#include <unordered_map>

namespace vk
{
	std::unordered_map<u64, std::vector<std::unique_ptr<vk::framebuffer_holder>>> g_framebuffers_cache;

	union framebuffer_storage_key
	{
		u64 encoded;

		struct
		{
			u64 width  : 16;   // Width of FBO
			u64 height : 16;   // Height of FBO
			u64 ia_ref : 1;    // Input attachment references?
		};

		framebuffer_storage_key(u16 width_, u16 height_, VkBool32 has_input_attachments)
			: width(width_), height(height_), ia_ref(has_input_attachments)
		{}
	};

	vk::framebuffer_holder* get_framebuffer(VkDevice dev, u16 width, u16 height, VkBool32 has_input_attachments, VkRenderPass renderpass, const std::vector<vk::image*>& image_list)
	{
		framebuffer_storage_key key(width, height, has_input_attachments);
		auto &queue = g_framebuffers_cache[key.encoded];

		for (auto &fbo : queue)
		{
			if (fbo->matches(image_list, width, height))
			{
				return fbo.get();
			}
		}

		std::vector<std::unique_ptr<vk::image_view>> image_views;
		image_views.reserve(image_list.size());

		for (auto &e : image_list)
		{
			const VkImageSubresourceRange subres = { e->aspect(), 0, 1, 0, 1 };
			image_views.push_back(std::make_unique<vk::image_view>(dev, e, VK_IMAGE_VIEW_TYPE_2D, vk::default_component_map, subres));
		}

		auto value = std::make_unique<vk::framebuffer_holder>(dev, renderpass, width, height, std::move(image_views));
		auto ret = value.get();

		queue.push_back(std::move(value));
		return ret;
	}

	vk::framebuffer_holder* get_framebuffer(VkDevice dev, u16 width, u16 height, VkBool32 has_input_attachments, VkRenderPass renderpass, VkFormat format, VkImage attachment)
	{
		framebuffer_storage_key key(width, height, has_input_attachments);
		auto &queue = g_framebuffers_cache[key.encoded];

		for (const auto &e : queue)
		{
			if (e->attachments[0]->info.image == attachment)
			{
				return e.get();
			}
		}

		VkImageSubresourceRange range = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 };
		std::vector<std::unique_ptr<vk::image_view>> views;

		views.push_back(std::make_unique<vk::image_view>(dev, attachment, VK_IMAGE_VIEW_TYPE_2D, format, vk::default_component_map, range));
		auto value = std::make_unique<vk::framebuffer_holder>(dev, renderpass, width, height, std::move(views));
		auto ret = value.get();

		queue.push_back(std::move(value));
		return ret;
	}

	void remove_unused_framebuffers()
	{
		// Remove stale framebuffers. Ref counted to prevent use-after-free
		for (auto It = g_framebuffers_cache.begin(); It != g_framebuffers_cache.end();)
		{
			It->second.erase(
				std::remove_if(It->second.begin(), It->second.end(), [](const auto& fbo)
				{
					return (fbo->unused_check_count() >= 2);
				}),
				It->second.end()
			);

			if (It->second.empty())
			{
				It = g_framebuffers_cache.erase(It);
			}
			else
			{
				++It;
			}
		}
	}

	void clear_framebuffer_cache()
	{
		g_framebuffers_cache.clear();
	}
}