#include "vulkan/vulkan.h"
#include <cassert>

import Crafter.Graphics;
using namespace Crafter;
import std;
import Crafter.Event;
import Crafter.Math;


int main() {
    Device::Initialize();
    Window window(1280, 720, "HelloVulkan");
    VkCommandBuffer cmd = window.StartInit();
    DescriptorHeapVulkan descriptorHeap;
    descriptorHeap.Initialize(1,2,0);

    VkSpecializationMapEntry entry = {
        .constantID = 0,
        .offset = 0,
        .size = sizeof(uint16_t)
    };

    VkSpecializationInfo specilizationInfo = {
        .mapEntryCount = 1,
        .pMapEntries = &entry,
        .dataSize = sizeof(uint16_t),
        .pData = &descriptorHeap.bufferStartElement
    };

    std::array<VulkanShader, 1> shaders{{
        {"raygen.spv", "main", VK_SHADER_STAGE_RAYGEN_BIT_KHR, &specilizationInfo}
    }};

    ShaderBindingTableVulkan shaderTable;
    shaderTable.Init(shaders);

    std::array<VkRayTracingShaderGroupCreateInfoKHR, 1> raygenGroups {{
        {
            .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
            .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR,
            .generalShader = 0,
            .closestHitShader = VK_SHADER_UNUSED_KHR,
            .anyHitShader = VK_SHADER_UNUSED_KHR,
            .intersectionShader = VK_SHADER_UNUSED_KHR,
        },
    }};
    std::array<VkRayTracingShaderGroupCreateInfoKHR, 0> missGroups;
    std::array<VkRayTracingShaderGroupCreateInfoKHR, 0> hitGroups;


    PipelineRTVulkan pipeline;
    pipeline.Init(cmd, raygenGroups, missGroups, hitGroups, shaderTable);

    window.FinishInit();

    RendertargetBase<3> rendertargetBase(1280, 720);
    RenderingElement2DVulkan<true, true> element(
        {
            0.5, //anchorX: relative position where this elements x anchor (top-left) is placed to its parent x anchor 
            0.5, //anchorY: relative position where this elements y anchor (top-left) is placed to its parent y anchor 
            0.5, //relativeSizeX: the relative x size this element should be scaled to compared to its parent
            0.5, //relativeSizeY: the relative y size this element should be scaled to compared to its parent
            0.5, //anchorOffsetX: the amount this element's anchor should be offset from the top left corner (0.5 to in the middle)
            0.5, //anchorOffsetY: the amount this element's anchor should be offset from the top left corner (0.5 to place it in the middle)
            0  //z: this elements Z position
        },
        2,
        1
    );
    for(std::uint8_t i = 0; i < Window::numFrames; i++) {
        element.buffers[i]->value[0] = {1, 0, 0, 1};
        element.buffers[i]->value[1] = {0, 1, 0, 1};
        element.buffers[i]->FlushDevice();
    }

    InitializeRenderingElement2DVulkanBuffer();
    element.UpdatePosition(rendertargetBase, rendertargetBase.transform);

    renderingElement2DVulkanTransformBuffer[0].FlushDevice();
    renderingElement2DVulkanTransformBuffer[1].FlushDevice();
    renderingElement2DVulkanTransformBuffer[2].FlushDevice();

    VkImageDescriptorInfoEXT imageInfo0 = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_DESCRIPTOR_INFO_EXT,
        .pView = &window.imageViews[0],
        .layout = VK_IMAGE_LAYOUT_GENERAL
    };

    VkImageDescriptorInfoEXT imageInfo1 = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_DESCRIPTOR_INFO_EXT,
        .pView = &window.imageViews[1],
        .layout = VK_IMAGE_LAYOUT_GENERAL
    };

    VkImageDescriptorInfoEXT imageInfo2 = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_DESCRIPTOR_INFO_EXT,
        .pView = &window.imageViews[2],
        .layout = VK_IMAGE_LAYOUT_GENERAL
    };
    
    VkDeviceAddressRangeKHR transformRange0 = {
        .address = renderingElement2DVulkanTransformBuffer[0].address,
        .size = renderingElement2DVulkanTransformBuffer[0].descriptor.range
    };

    VkDeviceAddressRangeKHR transformRange1 = {
        .address = renderingElement2DVulkanTransformBuffer[1].address,
        .size = renderingElement2DVulkanTransformBuffer[1].descriptor.range
    };

    VkDeviceAddressRangeKHR transformRange2 = {
        .address = renderingElement2DVulkanTransformBuffer[2].address,
        .size = renderingElement2DVulkanTransformBuffer[2].descriptor.range
    };

    VkDeviceAddressRangeKHR elRange0 = {
        .address = element.buffers[0]->address,
        .size = element.buffers[0]->descriptor.range
    };

    VkDeviceAddressRangeKHR elRange1 = {
        .address = element.buffers[1]->address,
        .size = element.buffers[1]->descriptor.range
    };

    VkDeviceAddressRangeKHR elRange2 = {
        .address = element.buffers[2]->address,
        .size = element.buffers[2]->descriptor.range
    };

    VkResourceDescriptorInfoEXT resources[9] = {
        {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
            .data = { .pImage = &imageInfo0 }
        },
        {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
            .data = { .pImage = &imageInfo1 }
        },
        {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
            .data = { .pImage = &imageInfo2 }
        },
        {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            .data = { .pAddressRange =  &transformRange0}
        },
        {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            .data = { .pAddressRange =  &transformRange1}
        },
        {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            .data = { .pAddressRange =  &transformRange2}
        },
        {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            .data = { .pAddressRange =  &elRange0}
        },
        {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            .data = { .pAddressRange =  &elRange1}
        },
        {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            .data = { .pAddressRange =  &elRange2}
        },
    };
    
    VkHostAddressRangeEXT destinations[9] = {
        {
            .address = descriptorHeap.resourceHeap[0].value,
            .size = Device::descriptorHeapProperties.imageDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[1].value,
            .size = Device::descriptorHeapProperties.imageDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[2].value,
            .size = Device::descriptorHeapProperties.imageDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[0].value + descriptorHeap.bufferStartOffset,
            .size = Device::descriptorHeapProperties.bufferDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[1].value + descriptorHeap.bufferStartOffset,
            .size = Device::descriptorHeapProperties.bufferDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[2].value + descriptorHeap.bufferStartOffset,
            .size = Device::descriptorHeapProperties.bufferDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[0].value + descriptorHeap.bufferStartOffset + Device::descriptorHeapProperties.bufferDescriptorSize * (renderingElement2DVulkans.size()),
            .size = Device::descriptorHeapProperties.bufferDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[1].value + descriptorHeap.bufferStartOffset + Device::descriptorHeapProperties.bufferDescriptorSize * (renderingElement2DVulkans.size()),
            .size = Device::descriptorHeapProperties.bufferDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[2].value + descriptorHeap.bufferStartOffset + Device::descriptorHeapProperties.bufferDescriptorSize * (renderingElement2DVulkans.size()),
            .size = Device::descriptorHeapProperties.bufferDescriptorSize
        },
    };

    Device::vkWriteResourceDescriptorsEXT(Device::device, 9, resources, destinations);

    descriptorHeap.resourceHeap[0].FlushDevice();
    descriptorHeap.resourceHeap[1].FlushDevice();
    descriptorHeap.resourceHeap[2].FlushDevice();

    window.pipeline = &pipeline;
    window.descriptorHeap = &descriptorHeap;

    window.Render();
    window.StartSync();
}