module;

#include <cstdint>
#include <string>
#include <vulkan/vulkan.h>
#include <vulkan/vulkan_wayland.h>
#include <vector>
#include <wayland-client.h>
#include <thread>
#include <iostream>

module Crafter.Graphics;
import Crafter.Event;
using namespace Crafter;

void WindowWaylandVulkan::CreateSwapchain()
{
    VkResult result;
    
    VkSurfaceCapabilitiesKHR capabilities;
    VulkanDevice::CHECK_VK_RESULT(vkGetPhysicalDeviceSurfaceCapabilitiesKHR(VulkanDevice::physDevice, vulkanSurface, &capabilities));

    uint32_t formatCount;
    VulkanDevice::CHECK_VK_RESULT(vkGetPhysicalDeviceSurfaceFormatsKHR(VulkanDevice::physDevice, vulkanSurface, &formatCount, NULL));

    std::vector<VkSurfaceFormatKHR> formats(formatCount);
    VulkanDevice::CHECK_VK_RESULT(vkGetPhysicalDeviceSurfaceFormatsKHR(VulkanDevice::physDevice, vulkanSurface, &formatCount, formats.data()));

    VkSurfaceFormatKHR chosenFormat = formats[0];

    for (uint32_t i = 0; i < formatCount; i++)
    {
        if (formats[i].format == VK_FORMAT_B8G8R8A8_UNORM)
        {
            chosenFormat = formats[i];
            break;
        }
    }

    format = chosenFormat.format;

    imageCount = capabilities.minImageCount + 1 < capabilities.maxImageCount ? capabilities.minImageCount + 1 : capabilities.minImageCount;

    VkSwapchainCreateInfoKHR swapchainCreateInfo = {};
    swapchainCreateInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
    swapchainCreateInfo.surface = vulkanSurface;
    swapchainCreateInfo.minImageCount = imageCount;
    swapchainCreateInfo.imageFormat = chosenFormat.format;
    swapchainCreateInfo.imageColorSpace = chosenFormat.colorSpace;
    swapchainCreateInfo.imageExtent.width = width;
    swapchainCreateInfo.imageExtent.height = height;
    swapchainCreateInfo.imageArrayLayers = 1;
    swapchainCreateInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
    swapchainCreateInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
    swapchainCreateInfo.preTransform = capabilities.currentTransform;
    swapchainCreateInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
    swapchainCreateInfo.presentMode = VK_PRESENT_MODE_MAILBOX_KHR;
    swapchainCreateInfo.clipped = 1;
    VulkanDevice::CHECK_VK_RESULT(vkCreateSwapchainKHR(VulkanDevice::device, &swapchainCreateInfo, NULL, &swapchain));
    
    VkAttachmentDescription attachment = {};
    attachment.format = format;
    attachment.samples = VK_SAMPLE_COUNT_1_BIT;
    attachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
    attachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
    attachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
    attachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
    attachment.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
    attachment.finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;

    VkAttachmentReference attachmentRef = {};
    attachmentRef.attachment = 0;
    attachmentRef.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

    VkSubpassDescription subpass = {};
    subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
    subpass.colorAttachmentCount = 1;
    subpass.pColorAttachments = &attachmentRef;

    VkRenderPassCreateInfo renderPasscreateInfo = {};
    renderPasscreateInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
    renderPasscreateInfo.flags = 0;
    renderPasscreateInfo.attachmentCount = 1;
    renderPasscreateInfo.pAttachments = &attachment;
    renderPasscreateInfo.subpassCount = 1;
    renderPasscreateInfo.pSubpasses = &subpass;
    VulkanDevice::CHECK_VK_RESULT(vkCreateRenderPass(VulkanDevice::device, &renderPasscreateInfo, NULL, &renderPass));
    
    VulkanDevice::CHECK_VK_RESULT(vkGetSwapchainImagesKHR(VulkanDevice::device, swapchain, &imageCount, NULL));

    std::vector<VkImage> images(imageCount);
    VulkanDevice::CHECK_VK_RESULT(vkGetSwapchainImagesKHR(VulkanDevice::device, swapchain, &imageCount, images.data()));

    swapchainElements.resize(imageCount);

    for (uint32_t i = 0; i < imageCount; i++)
    {
        VkCommandBufferAllocateInfo commandBufferAllocInfo = {};
        commandBufferAllocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
        commandBufferAllocInfo.commandPool = VulkanDevice::commandPool;
        commandBufferAllocInfo.commandBufferCount = 1;
        commandBufferAllocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
        vkAllocateCommandBuffers(VulkanDevice::device, &commandBufferAllocInfo, &swapchainElements[i].commandBuffer);
        
        swapchainElements[i].image = images[i];

        VkImageViewCreateInfo imageViewcreateInfo = {};
        imageViewcreateInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
        imageViewcreateInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
        imageViewcreateInfo.components.r = VK_COMPONENT_SWIZZLE_IDENTITY;
        imageViewcreateInfo.components.g = VK_COMPONENT_SWIZZLE_IDENTITY;
        imageViewcreateInfo.components.b = VK_COMPONENT_SWIZZLE_IDENTITY;
        imageViewcreateInfo.components.a = VK_COMPONENT_SWIZZLE_IDENTITY;
        imageViewcreateInfo.subresourceRange.baseMipLevel = 0;
        imageViewcreateInfo.subresourceRange.levelCount = 1;
        imageViewcreateInfo.subresourceRange.baseArrayLayer = 0;
        imageViewcreateInfo.subresourceRange.layerCount = 1;
        imageViewcreateInfo.image = swapchainElements[i].image;
        imageViewcreateInfo.format = format;
        imageViewcreateInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        VulkanDevice::CHECK_VK_RESULT(vkCreateImageView(VulkanDevice::device, &imageViewcreateInfo, NULL, &swapchainElements[i].imageView));
        
        VkFramebufferCreateInfo framebufferCreateInfo = {};
        framebufferCreateInfo.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
        framebufferCreateInfo.renderPass = renderPass;
        framebufferCreateInfo.attachmentCount = 1;
        framebufferCreateInfo.pAttachments = &swapchainElements[i].imageView;
        framebufferCreateInfo.width = width;
        framebufferCreateInfo.height = height;
        framebufferCreateInfo.layers = 1;
        VulkanDevice::CHECK_VK_RESULT(vkCreateFramebuffer(VulkanDevice::device, &framebufferCreateInfo, NULL, &swapchainElements[i].framebuffer));
    
        VkSemaphoreCreateInfo startSemaphoreCreateInfo = {};
        startSemaphoreCreateInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;
        VulkanDevice::CHECK_VK_RESULT(vkCreateSemaphore(VulkanDevice::device, &startSemaphoreCreateInfo, NULL, &swapchainElements[i].startSemaphore));
        
        VkSemaphoreCreateInfo endSemaphoreCreateInfo = {};
        endSemaphoreCreateInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;
        VulkanDevice::CHECK_VK_RESULT(vkCreateSemaphore(VulkanDevice::device, &endSemaphoreCreateInfo, NULL, &swapchainElements[i].endSemaphore));
   
        VkFenceCreateInfo fenceCreateInfo = {};
        fenceCreateInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
        fenceCreateInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;
        VulkanDevice::CHECK_VK_RESULT(vkCreateFence(VulkanDevice::device, &fenceCreateInfo, NULL, &swapchainElements[i].fence));

        swapchainElements[i].lastFence = VK_NULL_HANDLE;
    }
}

void WindowWaylandVulkan::DestroySwapchain()
{
    for (uint32_t i = 0; i < imageCount; i++)
    {
        vkDestroyFence(VulkanDevice::device, swapchainElements[i].fence, NULL);
        vkDestroySemaphore(VulkanDevice::device, swapchainElements[i].endSemaphore, NULL);
        vkDestroySemaphore(VulkanDevice::device, swapchainElements[i].startSemaphore, NULL);
        vkDestroyFramebuffer(VulkanDevice::device, swapchainElements[i].framebuffer, NULL);
        vkDestroyImageView(VulkanDevice::device, swapchainElements[i].imageView, NULL);
        vkFreeCommandBuffers(VulkanDevice::device, VulkanDevice::commandPool, 1, &swapchainElements[i].commandBuffer);
    }

    vkDestroyRenderPass(VulkanDevice::device, renderPass, NULL);
    vkDestroySwapchainKHR(VulkanDevice::device, swapchain, NULL);
}

PFN_vkCmdDrawMeshTasksEXT command;

WindowWaylandVulkan::WindowWaylandVulkan(std::string name, std::uint32_t width, std::uint32_t height) : WindowWayland(name, width, height) {
    VkWaylandSurfaceCreateInfoKHR createInfo = {};
    createInfo.sType = VK_STRUCTURE_TYPE_WAYLAND_SURFACE_CREATE_INFO_KHR;
    createInfo.display = display;
    createInfo.surface = surface;
    VulkanDevice::CHECK_VK_RESULT(vkCreateWaylandSurfaceKHR(VulkanDevice::instance, &createInfo, NULL, &vulkanSurface));
    CreateSwapchain();
    command = reinterpret_cast<PFN_vkCmdDrawMeshTasksEXT>(vkGetDeviceProcAddr(VulkanDevice::device, "vkCmdDrawMeshTasksEXT"));
}

void WindowWaylandVulkan::Start() {
    thread = std::thread([this](){
		while (open && wl_display_dispatch(display) != -1) {

        }
    });
    while(true) {
        SwapchainElement* currentElement = &swapchainElements[currentFrame];

        VulkanDevice::CHECK_VK_RESULT(vkWaitForFences(VulkanDevice::device, 1, &currentElement->fence, 1, UINT64_MAX));
        VulkanDevice::CHECK_VK_RESULT(vkAcquireNextImageKHR(VulkanDevice::device, swapchain, UINT64_MAX, currentElement->startSemaphore, NULL, &imageIndex));

        SwapchainElement* element = &swapchainElements[imageIndex];
        if (element->lastFence) {
            VulkanDevice::CHECK_VK_RESULT(vkWaitForFences(VulkanDevice::device, 1, &element->lastFence, 1, UINT64_MAX));
        }

        element->lastFence = currentElement->fence;

        VulkanDevice::CHECK_VK_RESULT(vkResetFences(VulkanDevice::device, 1, &currentElement->fence));

        VkCommandBufferBeginInfo commandBeginInfo = {};
        commandBeginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
        commandBeginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;

        VulkanDevice::CHECK_VK_RESULT(vkBeginCommandBuffer(element->commandBuffer, &commandBeginInfo));
        
        VkClearValue clearValue = {{
            1.0f,
            0.0f,
            1.0f,
            1.0f
        }};

        VkRenderPassBeginInfo renderBeginInfo = {};
        renderBeginInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
        renderBeginInfo.renderPass = renderPass;
        renderBeginInfo.framebuffer = element->framebuffer;
        renderBeginInfo.renderArea.offset.x = 0;
        renderBeginInfo.renderArea.offset.y = 0;
        renderBeginInfo.renderArea.extent.width = width;
        renderBeginInfo.renderArea.extent.height = height;
        renderBeginInfo.clearValueCount = 1;
        renderBeginInfo.pClearValues = &clearValue;

        vkCmdBeginRenderPass(element->commandBuffer, &renderBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

        VkViewport viewport = {0, 0, static_cast<float>(width), static_cast<float>(height), 0, 1};
		vkCmdSetViewport(element->commandBuffer, 0, 1, &viewport);

        VkRect2D scissor = {{static_cast<std::int32_t>(width), static_cast<std::int32_t>(height)},{0,0}};
		vkCmdSetScissor(element->commandBuffer, 0, 1, &scissor);

        vkCmdBindDescriptorSets(element->commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, VulkanPipeline<VulkanShader<"test.spirv", "main", VK_SHADER_STAGE_MESH_BIT_EXT>, VulkanShader<"test2.spirv", "main", VK_SHADER_STAGE_FRAGMENT_BIT>>::layout, 0, 1, &VulkanPipeline<VulkanShader<"test.spirv", "main", VK_SHADER_STAGE_MESH_BIT_EXT>, VulkanShader<"test2.spirv", "main", VK_SHADER_STAGE_FRAGMENT_BIT>>::descriptor_set, 0, nullptr);
        vkCmdBindPipeline(element->commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, VulkanPipeline<VulkanShader<"test.spirv", "main", VK_SHADER_STAGE_MESH_BIT_EXT>, VulkanShader<"test2.spirv", "main", VK_SHADER_STAGE_FRAGMENT_BIT>>::pipeline);
        
        command(element->commandBuffer, 1, 1, 1);

        vkCmdEndRenderPass(element->commandBuffer);

        VulkanDevice::CHECK_VK_RESULT(vkEndCommandBuffer(element->commandBuffer));

        const VkPipelineStageFlags waitStage = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;

        VkSubmitInfo submitInfo = {};
        submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
        submitInfo.waitSemaphoreCount = 1;
        submitInfo.pWaitSemaphores = &currentElement->startSemaphore;
        submitInfo.pWaitDstStageMask = &waitStage;
        submitInfo.commandBufferCount = 1;
        submitInfo.pCommandBuffers = &element->commandBuffer;
        submitInfo.signalSemaphoreCount = 1;
        submitInfo.pSignalSemaphores = &currentElement->endSemaphore;

        VulkanDevice::CHECK_VK_RESULT(vkQueueSubmit(VulkanDevice::queue, 1, &submitInfo, currentElement->fence));

        VkPresentInfoKHR presentInfo = {};
        presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
        presentInfo.waitSemaphoreCount = 1;
        presentInfo.pWaitSemaphores = &currentElement->endSemaphore;
        presentInfo.swapchainCount = 1;
        presentInfo.pSwapchains = &swapchain;
        presentInfo.pImageIndices = &imageIndex;

        VulkanDevice::CHECK_VK_RESULT(vkQueuePresentKHR(VulkanDevice::queue, &presentInfo));

        currentFrame = (currentFrame + 1) % imageCount;

        wl_display_roundtrip(display);
    }
}