#include <vulkan/vulkan.h>

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

typedef VulkanShader<"raygen.spv", "main", VK_SHADER_STAGE_RAYGEN_BIT_KHR> Raygenspv;
typedef VulkanShader<"closesthit.spv", "main", VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR> Closesthitspv;
typedef VulkanShader<"miss.spv", "main", VK_SHADER_STAGE_MISS_BIT_KHR> Misspv;
typedef std::tuple<Raygenspv, Misspv, Closesthitspv> AllShaders;
typedef std::tuple<
    ShaderGroup<0, VK_SHADER_UNUSED_KHR, VK_SHADER_UNUSED_KHR, VK_SHADER_UNUSED_KHR>,
    ShaderGroup<1, VK_SHADER_UNUSED_KHR, VK_SHADER_UNUSED_KHR, VK_SHADER_UNUSED_KHR>,
    ShaderGroup<VK_SHADER_UNUSED_KHR, 2, VK_SHADER_UNUSED_KHR, VK_SHADER_UNUSED_KHR>
> ShaderGroups;
typedef PipelineRTVulkan<AllShaders, ShaderGroups> Pipeline;
typedef DescriptorSetLayoutVulkan<1, {{
    {
        .binding = 0,
        .descriptorType = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR,
        .descriptorCount = 1,
        .stageFlags = VK_SHADER_STAGE_RAYGEN_BIT_KHR,
    },
}}> descriptorSetLayoutTlas;
typedef DescriptorSetLayoutVulkan<1, {{
    {
        .binding = 0,
        .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
        .descriptorCount = 1,
        .stageFlags = VK_SHADER_STAGE_RAYGEN_BIT_KHR,
    }
}}> descriptorSetLayoutImage;

int main() {
    VulkanDevice::CreateDevice();
    WindowVulkan window(1280, 720, "HelloVulkan");
    VkCommandBuffer cmd = window.StartInit();

    Raygenspv::CreateShader();
    Closesthitspv::CreateShader();
    Misspv::CreateShader();

    ShaderBindingTableVulkan<AllShaders>::Init();

    descriptorSetLayoutTlas::Init();
    descriptorSetLayoutImage::Init();
    std::array<VkDescriptorSetLayout, 4> layouts {{descriptorSetLayoutTlas::layout, descriptorSetLayoutImage::layout, descriptorSetLayoutImage::layout, descriptorSetLayoutImage::layout}};

    DescriptorPool pool;
    pool.sets.resize(4);
    pool.BuildPool(DescriptorPool::GetPoolSizes<descriptorSetLayoutTlas, descriptorSetLayoutImage, descriptorSetLayoutImage, descriptorSetLayoutImage>(), layouts);

    Pipeline::Init(cmd, layouts);
    window.SetPipelineRT<Pipeline>();

    Mesh triangleMesh;
    std::array<Vertex, 3> verts {{{-150, -150, 100}, {0, 150, 100}, {150, -150, 100}}};
    std::array<std::uint32_t, 3> index {{2,1,0}};
    triangleMesh.Build(verts, index, cmd);
    RenderingElement3DVulkan& el = RenderingElement3DVulkan::elements.emplace_back(triangleMesh);
    MatrixRowMajor<float, 4, 3, 1> transform = MatrixRowMajor<float, 4, 3, 1>::Identity();
    std::memcpy(el.instance.transform.matrix, transform.m, sizeof(transform.m));

    RenderingElement3DVulkan::tlases.resize(WindowVulkan::numFrames);
    RenderingElement3DVulkan::BuildTLAS(cmd, 0);
    RenderingElement3DVulkan::BuildTLAS(cmd, 1);
    RenderingElement3DVulkan::BuildTLAS(cmd, 2);
    
    VkDescriptorImageInfo imageInfo[WindowVulkan::numFrames] = {  
        {
            .imageView = window.imageViews[0],
            .imageLayout = VK_IMAGE_LAYOUT_GENERAL
        },
        {
            .imageView = window.imageViews[1],
            .imageLayout = VK_IMAGE_LAYOUT_GENERAL
        },
        {
            .imageView = window.imageViews[2],
            .imageLayout = VK_IMAGE_LAYOUT_GENERAL
        },
    };

    VkWriteDescriptorSetAccelerationStructureKHR writeDescriptorSetAccelerationStructure {
        .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET_ACCELERATION_STRUCTURE_KHR,
        .accelerationStructureCount = 1,
        .pAccelerationStructures = &RenderingElement3DVulkan::tlases[0].accelerationStructure
    };

    VkWriteDescriptorSet write[4];

    write[0] = {
        .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .pNext = &writeDescriptorSetAccelerationStructure,
        .dstSet = pool.sets[0],
        .dstBinding = 0,
        .dstArrayElement = 0,
        .descriptorCount = 1,
        .descriptorType = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR,
    };

    for(std::uint32_t i = 0; i < WindowVulkan::numFrames; i++) {
        write[i+1] = {
            .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
            .dstSet = pool.sets[i+1],
            .dstBinding = 0,
            .dstArrayElement = 0,
            .descriptorCount = 1,
            .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
            .pImageInfo = &imageInfo[i]
        };
    }
    vkUpdateDescriptorSets(VulkanDevice::device, 4, write, 0, nullptr);
    window.descriptorsRt = {pool.sets[0], pool.sets[1]};

    /*
        FinishInit executes all commands recorded to StartInit.
        This must be called before the the event loops starts if you called StartInit before.
    */
    window.FinishInit();

    Animation<std::tuple<float>> anim({
        {std::chrono::seconds(3), -600, 600},
    });

	anim.Start(std::chrono::high_resolution_clock::now());

    EventListener<void> updateListener(&window.onRender, [&](){
        float value = std::get<0>(anim.Play(window.currentFrameTime.now));

        if(anim.currentFrame == anim.keyframes.size()) {
            anim.Start(window.currentFrameTime.now);
        }

        MatrixRowMajor<float, 4, 3, 1> transform = MatrixRowMajor<float, 4, 3, 1>::Translation(value, 0, 0);
        std::memcpy(el.instance.transform.matrix, transform.m, sizeof(transform.m));
        RenderingElement3DVulkan::BuildTLAS(window.drawCmdBuffers[window.currentBuffer], window.currentBuffer);

        VkWriteDescriptorSetAccelerationStructureKHR writeDescriptorSetAccelerationStructure {
            .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET_ACCELERATION_STRUCTURE_KHR,
            .accelerationStructureCount = 1,
            .pAccelerationStructures = &RenderingElement3DVulkan::tlases[window.currentBuffer].accelerationStructure
        };

        VkWriteDescriptorSet write = {
            .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
            .pNext = &writeDescriptorSetAccelerationStructure,
            .dstSet = pool.sets[0],
            .dstBinding = 0,
            .dstArrayElement = 0,
            .descriptorCount = 1,
            .descriptorType = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR,
        };
        vkUpdateDescriptorSets(VulkanDevice::device, 1, &write, 0, nullptr);

        window.descriptorsRt[1] = pool.sets[window.currentBuffer+1];
    });
    
    window.Render();
	window.StartUpdate();
    window.StartSync();
}