/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#include "DepthBufferQuery.hpp"
#include "Renderer.hpp"
#include <vulkan/vulkan.hpp>

//TODO resize

namespace OpenVulkano::Vulkan
{
	namespace
	{
		constexpr uint32_t SAMPLE_SIZE_WIDTH = 5;
		constexpr uint32_t SAMPLE_SIZE_HEIGHT = 5;
		constexpr uint32_t SAMPLE_SIZE = SAMPLE_SIZE_WIDTH * SAMPLE_SIZE_HEIGHT;
		constexpr uint32_t SAMPLE_CENTER = SAMPLE_SIZE / 2;
	}

	void DepthBufferQuery::Init()
	{
		auto device = renderer.GetContext().device->device;

		vk::BufferCreateInfo bufferInfo = { {}, SAMPLE_SIZE * sizeof(float), vk::BufferUsageFlagBits::eTransferDst };
		bufferDepth = device.createBuffer(bufferInfo);
		const vk::MemoryRequirements memRequirements = device.getBufferMemoryRequirements(bufferDepth);
		size_t size = memRequirements.size;
		const vk::MemoryAllocateInfo memAllocInfo = { size, renderer.GetContext().device->GetMemoryType(memRequirements.memoryTypeBits, vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostVisible) };

		memory = device.allocateMemory(memAllocInfo);
		cpuDepthBuffer = static_cast<float*>(device.mapMemory(memory, 0, VK_WHOLE_SIZE, {}));

		device.bindBufferMemory(bufferDepth, memory, 0);
	}

	void DepthBufferQuery::Close()
	{
		auto device = renderer.GetContext().device->device;
		device.destroy(bufferDepth);
		device.unmapMemory(memory);
		device.free(memory);
		cpuDepthBuffer = nullptr;
	}
	
	void DepthBufferQuery::Resize(uint32_t width, uint32_t height)
	{
		//TODO
	}
	
	float DepthBufferQuery::GetQueriedValue() const
	{
		if (cpuDepthBuffer[0] == -2) return -2;
		if (cpuDepthBuffer[SAMPLE_CENTER] > 0 && cpuDepthBuffer[SAMPLE_CENTER] < 1) return cpuDepthBuffer[SAMPLE_CENTER];
		std::sort(cpuDepthBuffer, cpuDepthBuffer + SAMPLE_SIZE);
		uint32_t start = UINT32_MAX, end = UINT32_MAX;
		for (uint32_t i = 0; i < SAMPLE_SIZE; i++)
		{
			float f = cpuDepthBuffer[i];
			if (f > 0 && start == UINT32_MAX) start = i;
			if (f < 1 && end == UINT32_MAX) end = i;
			else if (f >= 1) break;
		}
		if (start == UINT32_MAX || end == UINT32_MAX) return cpuDepthBuffer[SAMPLE_CENTER];
		return cpuDepthBuffer[(start + end) / 2];
	}
	
	vk::Offset3D DepthBufferQuery::GetCopyOffset() const
	{
		vk::Extent3D depthExtent = renderer.GetContext().swapChain.GetCurrentDepthBuffer().extent;
		int32_t x = static_cast<int32_t>(depthQueryCoordinates.x * depthExtent.width);
		x = std::min<int32_t>(depthExtent.width - SAMPLE_SIZE_WIDTH, std::max(0, x));
		int32_t y = static_cast<int32_t>(depthQueryCoordinates.y * depthExtent.height);
		y = std::min<int32_t>(depthExtent.height - SAMPLE_SIZE_HEIGHT, std::max(0, y));
		return { x, y, 0 };
	}

	void DepthBufferQuery::Encode(vk::CommandBuffer& commandBuffer)
	{
		if (!copyDepthBuffer) return;
		copyDepthBuffer = false;
		std::fill(cpuDepthBuffer, cpuDepthBuffer + SAMPLE_SIZE, -2.0f); // Invalidate data in buffer to allow detecting if copy is done
		const vk::ImageAspectFlags aspectMask = vk::ImageAspectFlagBits::eDepth | vk::ImageAspectFlagBits::eStencil;
		const Image& depthBufferImage = renderer.GetContext().swapChain.GetCurrentDepthBuffer();

		constexpr vk::Extent3D copySize = { SAMPLE_SIZE_WIDTH, SAMPLE_SIZE_HEIGHT, 1 };
		const vk::ImageSubresourceLayers layout = { vk::ImageAspectFlagBits::eDepth, 0, 0, 1 };
		vk::BufferImageCopy imgCopy = { 0, 0, 0, layout, GetCopyOffset(), copySize };

		const vk::ImageMemoryBarrier imgMemBarrier({}, vk::AccessFlagBits::eTransferRead, vk::ImageLayout::eDepthStencilAttachmentOptimal, vk::ImageLayout::eTransferSrcOptimal, 0, 0, depthBufferImage.image, vk::ImageSubresourceRange(aspectMask, 0, 1, 0, 1));
		commandBuffer.pipelineBarrier(vk::PipelineStageFlagBits::eTopOfPipe, vk::PipelineStageFlagBits::eTransfer, {}, nullptr, nullptr, imgMemBarrier);

		commandBuffer.copyImageToBuffer(depthBufferImage.image, vk::ImageLayout::eTransferSrcOptimal, bufferDepth, 1, &imgCopy);

		depthBufferImage.SetLayout(commandBuffer, aspectMask, vk::ImageLayout::eDepthStencilAttachmentOptimal, vk::ImageLayout::eTransferSrcOptimal);
	}
}