// Blaze.cpp : This file contains the 'main' function. Program execution begins and ends there.
//


#include <array>
#include <cassert>
#include <span>

#include <volk.h>

#define SDL_MAIN_USE_CALLBACKS 1
#include <memory>
#include <SDL3/SDL.h>
#include <SDL3/SDL_main.h>
#include <SDL3/SDL_filesystem.h>
#include <SDL3/SDL_vulkan.h>

#include "Frame.h"
#include "MacroUtils.h"
#include "MathUtil.h"
#include "RenderDevice.h"

constexpr uint32_t WIDTH = 1280;
constexpr uint32_t HEIGHT = 720;
constexpr uint32_t NUM_FRAMES = 3;

struct MiscData
{
	VkPipelineLayout pipelineLayout;
	VkPipeline trianglePipeline;

	VkImageMemoryBarrier2 acquireToRenderBarrier;
	VkDependencyInfo acquireToRenderDependency;
	VkImageMemoryBarrier2 renderToPresentBarrier;
	VkDependencyInfo renderToPresentDependency;

	void init(RenderDevice const& renderDevice);
	void cleanup(RenderDevice const& renderDevice);
};

struct AppState
{
	SDL_Window* window;
	std::unique_ptr<RenderDevice> renderDevice;
	MiscData miscData;

	[[nodiscard]] bool isInit() const;
	void cleanup();

	AppState();

	AppState(AppState const& other) = delete;
	AppState(AppState&& other) noexcept = delete;

	AppState& operator=(AppState const& other) = delete;
	AppState& operator=(AppState&& other) noexcept = delete;

	~AppState();
};

static_assert(sizeof(AppState) < 1024);

bool AppState::isInit() const
{
	return window and renderDevice and renderDevice->isInit();
}

void AppState::cleanup()
{
	if (!isInit()) return;

	renderDevice->waitIdle();

	miscData.cleanup(*renderDevice);
	renderDevice->cleanup();
	SDL_DestroyWindow(window);
}

AppState::AppState()
{
	window = SDL_CreateWindow("Blaze Test", WIDTH, HEIGHT, SDL_WINDOW_VULKAN);
	if (!window)
	{
		SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "%s", SDL_GetError());
		abort();
	}

	renderDevice = std::make_unique<RenderDevice>(RenderDevice::CreateInfo{ .window = window });
	if (!renderDevice->isInit())
	{
		SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "RenderDevice failed to init");
		abort();
	}

	miscData.init(*renderDevice);
}

AppState::~AppState()
{
	cleanup();
}

SDL_AppResult SDL_AppInit(void** pAppState, int, char**)
{
	SDL_Init(SDL_INIT_VIDEO | SDL_INIT_EVENTS);

	*pAppState = new AppState{};
	AppState& appState = *static_cast<AppState*>(*pAppState);

	if (!appState.isInit())
	{
		return SDL_APP_FAILURE;
	}

	return SDL_APP_CONTINUE;
}

SDL_AppResult SDL_AppIterate(void* appstate)
{
	AppState& appState = *static_cast<AppState*>(appstate);
	RenderDevice& renderDevice = *appState.renderDevice;
	MiscData& misc = appState.miscData;
	Frame& currentFrame = renderDevice.frames[renderDevice.frameIndex];

	VK_CHECK(vkWaitForFences(renderDevice.device, 1, &currentFrame.frameReadyToReuse, VK_TRUE, std::numeric_limits<uint32_t>::max()));
	// All resources of frame 'frameIndex' are free.

	uint32_t currentImageIndex;
	VK_CHECK(vkAcquireNextImageKHR(renderDevice.device, renderDevice.swapchain, std::numeric_limits<uint32_t>::max(), currentFrame.imageAcquiredSemaphore, nullptr, &currentImageIndex));

	VK_CHECK(vkResetFences(renderDevice.device, 1, &currentFrame.frameReadyToReuse));
	VK_CHECK(vkResetCommandPool(renderDevice.device, currentFrame.commandPool, 0));

	misc.acquireToRenderBarrier.image = renderDevice.swapchainImages[currentImageIndex];
	misc.renderToPresentBarrier.image = renderDevice.swapchainImages[currentImageIndex];

	VkCommandBuffer cmd = currentFrame.commandBuffer;

	VkCommandBufferBeginInfo constexpr beginInfo = {
		.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
		.pNext = nullptr,
		.flags = 0,
		.pInheritanceInfo = nullptr,
	};

	VkClearColorValue constexpr static BLACK_CLEAR = {
		.float32 = { 0.0f, 0.0f, 0.0f, 1.0f },
	};

	VK_CHECK(vkBeginCommandBuffer(cmd, &beginInfo));
	{
		VkRenderingAttachmentInfo const attachmentInfo = {
			.sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO,
			.pNext = nullptr,
			.imageView = renderDevice.swapchainViews[currentImageIndex],
			.imageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
			.resolveMode = VK_RESOLVE_MODE_NONE,
			.resolveImageView = nullptr,
			.resolveImageLayout = VK_IMAGE_LAYOUT_UNDEFINED,
			.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
			.storeOp = VK_ATTACHMENT_STORE_OP_STORE,
			.clearValue = {.color = BLACK_CLEAR},
		};

		VkRenderingInfo renderingInfo = {
			.sType = VK_STRUCTURE_TYPE_RENDERING_INFO,
			.pNext = nullptr,
			.flags = 0,
			.renderArea = {.offset = {0,0}, .extent = renderDevice.swapchainExtent},
			.layerCount = 1,
			.viewMask = 0,
			.colorAttachmentCount = 1,
			.pColorAttachments = &attachmentInfo,
			.pDepthAttachment = nullptr,
			.pStencilAttachment = nullptr,
		};

		vkCmdPipelineBarrier2(cmd, &misc.acquireToRenderDependency);
		vkCmdBeginRendering(cmd, &renderingInfo);
		{
			VkViewport viewport = {
				.x = 0,
				.y = static_cast<float>(renderDevice.swapchainExtent.height),
				.width = static_cast<float>(renderDevice.swapchainExtent.width),
				.height = -static_cast<float>(renderDevice.swapchainExtent.height),
				.minDepth = 0.0f,
				.maxDepth = 1.0f,
			};
			vkCmdSetViewport(cmd, 0, 1, &viewport);
			VkRect2D scissor = {
				.offset = {0, 0},
				.extent = renderDevice.swapchainExtent,
			};
			vkCmdSetScissor(cmd, 0, 1, &scissor);

			// Render Something?
			vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS, misc.trianglePipeline);
			vkCmdDraw(cmd, 3, 1, 0, 0);

		}
		vkCmdEndRendering(cmd);
		vkCmdPipelineBarrier2(cmd, &misc.renderToPresentDependency);
	}
	VK_CHECK(vkEndCommandBuffer(cmd));

	VkPipelineStageFlags stageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
	VkSubmitInfo const submitInfo = {
		.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
		.pNext = nullptr,
		.waitSemaphoreCount = 1,
		.pWaitSemaphores = &currentFrame.imageAcquiredSemaphore,
		.pWaitDstStageMask = &stageMask,
		.commandBufferCount = 1,
		.pCommandBuffers = &cmd,
		.signalSemaphoreCount = 1,
		.pSignalSemaphores = &currentFrame.renderFinishedSemaphore,
	};
	VK_CHECK(vkQueueSubmit(renderDevice.directQueue, 1, &submitInfo, currentFrame.frameReadyToReuse));

	VkPresentInfoKHR const presentInfo = {
		.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
		.pNext = nullptr,
		.waitSemaphoreCount = 1,
		.pWaitSemaphores = &currentFrame.renderFinishedSemaphore,
		.swapchainCount = 1,
		.pSwapchains = &renderDevice.swapchain,
		.pImageIndices = &currentImageIndex,
		.pResults = nullptr,
	};

	VK_CHECK(vkQueuePresentKHR(renderDevice.directQueue, &presentInfo));

	renderDevice.frameIndex = (renderDevice.frameIndex + 1) % NUM_FRAMES;

	return SDL_APP_CONTINUE;
}

SDL_AppResult SDL_AppEvent(void*, SDL_Event* event)
{
	if (event->type == SDL_EVENT_QUIT)
	{
		return SDL_APP_SUCCESS;
	}

	return SDL_APP_CONTINUE;
}

void SDL_AppQuit(void* appstate, SDL_AppResult)
{
	AppState *appState = static_cast<AppState*> (appstate);

	appState->cleanup();

	// Not necessary, this memory will be winked out.
	delete appState;
}

void MiscData::init(RenderDevice const& renderDevice)
{
	VkDevice const device = renderDevice.device;
	{
		size_t dataSize;
		void* rawData = SDL_LoadFile("Triangle.spv", &dataSize);
		ASSERT(dataSize % 4 == 0);

		if (not rawData)
		{
			SDL_LogError(SDL_LOG_CATEGORY_SYSTEM, "%s", SDL_GetError());
			abort();
		}

		auto data = static_cast<uint32_t const*>(rawData);

		VkShaderModuleCreateInfo const shaderModuleCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.codeSize = dataSize,
			.pCode = data,
		};

		VkShaderModule shaderModule;
		VK_CHECK(vkCreateShaderModule(device, &shaderModuleCreateInfo, nullptr, &shaderModule));

		VkPipelineLayoutCreateInfo constexpr pipelineLayoutCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.setLayoutCount = 0,
			.pSetLayouts = nullptr,
			.pushConstantRangeCount = 0,
			.pPushConstantRanges = nullptr,
		};
		VK_CHECK(vkCreatePipelineLayout(device, &pipelineLayoutCreateInfo, nullptr, &pipelineLayout));

		std::array stages = {
			VkPipelineShaderStageCreateInfo{
				.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
				.pNext = nullptr,
				.flags = 0,
				.stage = VK_SHADER_STAGE_VERTEX_BIT,
				.module = shaderModule,
				.pName = "VertexMain",
				.pSpecializationInfo = nullptr,
			},
			VkPipelineShaderStageCreateInfo{
				.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
				.pNext = nullptr,
				.flags = 0,
				.stage = VK_SHADER_STAGE_FRAGMENT_BIT,
				.module = shaderModule,
				.pName = "FragmentMain",
				.pSpecializationInfo = nullptr,
			}
		};

		VkPipelineVertexInputStateCreateInfo constexpr vertexInputState = {
			.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.vertexBindingDescriptionCount = 0,
			.pVertexBindingDescriptions = nullptr,
			.vertexAttributeDescriptionCount = 0,
			.pVertexAttributeDescriptions = nullptr,
		};

		VkPipelineInputAssemblyStateCreateInfo constexpr inputAssembly = {
			.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST,
			.primitiveRestartEnable = VK_FALSE,
		};

		VkPipelineTessellationStateCreateInfo constexpr tessellationState = {
			.sType = VK_STRUCTURE_TYPE_PIPELINE_TESSELLATION_STATE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.patchControlPoints = 0,
		};

		VkPipelineViewportStateCreateInfo constexpr viewportState = {
			.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.viewportCount = 1,
			.pViewports = nullptr,
			.scissorCount = 1,
			.pScissors = nullptr,
		};

		VkPipelineRasterizationStateCreateInfo constexpr rasterizationState = {
			.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.depthClampEnable = VK_TRUE,
			.rasterizerDiscardEnable = VK_FALSE,
			.polygonMode = VK_POLYGON_MODE_FILL,
			.cullMode = VK_CULL_MODE_NONE,
			.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE,
			.depthBiasEnable = VK_FALSE,
			.depthBiasConstantFactor = 0.0f,
			.depthBiasClamp = 0.0f,
			.depthBiasSlopeFactor = 0.0f,
			.lineWidth = 1.0f,
		};

		VkPipelineMultisampleStateCreateInfo constexpr multisampleState = {
			.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT,
			.sampleShadingEnable = VK_FALSE,
			.minSampleShading = 0.0f,
			.pSampleMask = nullptr,
			.alphaToCoverageEnable = VK_FALSE,
			.alphaToOneEnable = VK_FALSE,
		};

		VkPipelineDepthStencilStateCreateInfo constexpr depthStencilState = {
			.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.depthTestEnable = VK_FALSE,
			.depthWriteEnable = VK_FALSE,
			.depthCompareOp = VK_COMPARE_OP_ALWAYS,
			.depthBoundsTestEnable = VK_FALSE,
			.stencilTestEnable = VK_FALSE,
			.front = {},
			.back = {},
			.minDepthBounds = 0.0f,
			.maxDepthBounds = 1.0f,
		};

		VkPipelineColorBlendAttachmentState constexpr colorBlendAttachmentState = {
			.blendEnable = VK_FALSE,
			.srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA,
			.dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA,
			.colorBlendOp = VK_BLEND_OP_ADD,
			.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE,
			.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO,
			.alphaBlendOp = VK_BLEND_OP_ADD,
			.colorWriteMask = VK_COLOR_COMPONENT_R_BIT
							| VK_COLOR_COMPONENT_G_BIT
							| VK_COLOR_COMPONENT_B_BIT
							| VK_COLOR_COMPONENT_A_BIT,
		};

		VkPipelineColorBlendStateCreateInfo const colorBlendState = {
			.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.logicOpEnable = VK_FALSE,
			.logicOp = VK_LOGIC_OP_COPY,
			.attachmentCount = 1,
			.pAttachments = &colorBlendAttachmentState,
			.blendConstants = {0.0f, 0.0f, 0.0f, 0.0f},
		};

		std::array constexpr dynamicStates = {
			VK_DYNAMIC_STATE_VIEWPORT,
			VK_DYNAMIC_STATE_SCISSOR
		};

		VkPipelineDynamicStateCreateInfo const dynamicStateCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.dynamicStateCount = static_cast<uint32_t>(dynamicStates.size()),
			.pDynamicStates = dynamicStates.data()
		};

		VkPipelineRenderingCreateInfoKHR const renderingCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO_KHR,
			.colorAttachmentCount = 1,
			.pColorAttachmentFormats = &renderDevice.swapchainFormat,
		};

		VkGraphicsPipelineCreateInfo const graphicsPipelineCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
			.pNext = &renderingCreateInfo,
			.flags = 0,
			.stageCount = static_cast<uint32_t>(stages.size()),
			.pStages = stages.data(),
			.pVertexInputState = &vertexInputState,
			.pInputAssemblyState = &inputAssembly,
			.pTessellationState = &tessellationState,
			.pViewportState = &viewportState,
			.pRasterizationState = &rasterizationState,
			.pMultisampleState = &multisampleState,
			.pDepthStencilState = &depthStencilState,
			.pColorBlendState = &colorBlendState,
			.pDynamicState = &dynamicStateCreateInfo,
			.layout = pipelineLayout,
			.renderPass = nullptr,
			.subpass = 0,
			.basePipelineHandle = nullptr,
			.basePipelineIndex = 0,
		};

		VK_CHECK(vkCreateGraphicsPipelines(device, nullptr, 1, &graphicsPipelineCreateInfo, nullptr, &trianglePipeline));

		vkDestroyShaderModule(device, shaderModule, nullptr);

		SDL_free(rawData);
	}

	acquireToRenderBarrier = {
		.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2,
		.pNext = nullptr,
		.srcStageMask = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
		.srcAccessMask = VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT,
		.dstStageMask = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
		.dstAccessMask = VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT,
		.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED,
		.newLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
		.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
		.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
		.subresourceRange = {
			.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
			.baseMipLevel = 0,
			.levelCount = 1,
			.baseArrayLayer = 0,
			.layerCount = 1,
		}
	};
	acquireToRenderDependency = {
		.sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
		.pNext = nullptr,
		.dependencyFlags = 0,
		.memoryBarrierCount = 0,
		.pMemoryBarriers = nullptr,
		.bufferMemoryBarrierCount = 0,
		.pBufferMemoryBarriers = nullptr,
		.imageMemoryBarrierCount = 1,
		.pImageMemoryBarriers = &acquireToRenderBarrier,
	};

	renderToPresentBarrier = {
		.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2,
		.pNext = nullptr,
		.srcStageMask = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
		.srcAccessMask = VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT,
		.dstStageMask = VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT,
		.dstAccessMask = VK_ACCESS_2_NONE,
		.oldLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
		.newLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
		.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
		.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
		.subresourceRange = {
			.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
			.baseMipLevel = 0,
			.levelCount = 1,
			.baseArrayLayer = 0,
			.layerCount = 1,
		}
	};
	renderToPresentDependency = {
		.sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
		.pNext = nullptr,
		.dependencyFlags = 0,
		.memoryBarrierCount = 0,
		.pMemoryBarriers = nullptr,
		.bufferMemoryBarrierCount = 0,
		.pBufferMemoryBarriers = nullptr,
		.imageMemoryBarrierCount = 1,
		.pImageMemoryBarriers = &renderToPresentBarrier,
	};

}

void MiscData::cleanup(RenderDevice const& renderDevice)
{
	VkDevice const device = renderDevice.device;

	vkDestroyPipeline(device, trianglePipeline, nullptr);
	vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
}