/*
 * 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
{
	void DepthBufferQuery::Init()
	{
		auto device = renderer.GetContext().device->device;

		vk::BufferCreateInfo bufferInfo = { {}, 25 * 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[12] > 0 && cpuDepthBuffer[12] < 1) return cpuDepthBuffer[12];
		double val = 0;
		int validCount = 0;
		for (int i = 0; i < 25; i++)
		{
			float f = cpuDepthBuffer[i];
			if (f > 0 && f < 1)
			{
				val += f;
				validCount++;
			}
		}
		if (validCount == 0) return 1; // Prevent divide by 0
		return val / validCount;
	}
	
	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 - 5, std::max(0, x));
		
		int32_t y = static_cast<int32_t>(depthQueryCoordinates.y * depthExtent.height);
		y = std::min<int32_t>(depthExtent.height - 5, std::max(0, y));
		
		return { x, y, 0 };
	}

	void DepthBufferQuery::Encode(vk::CommandBuffer& commandBuffer)
	{
		if (!copyDepthBuffer) return;
		copyDepthBuffer = false;
		std::fill(cpuDepthBuffer, cpuDepthBuffer + 25, -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 = { 5, 5, 1 };
		const vk::ImageSubresourceLayers layout = { vk::ImageAspectFlagBits::eDepth, 0, 0, 1 };
		vk::BufferImageCopy imgCopy = { 0, 5, 5, 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);
	}
}