#include "vulkan/vulkan.h"

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

typedef VulkanShaderConst<"raygen.spv", "main", VK_SHADER_STAGE_RAYGEN_BIT_KHR> Raygenspv;
typedef VulkanShaderConst<"closesthit.spv", "main", VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR> Closesthitspv;
typedef VulkanShaderConst<"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 PipelineRTVulkanConst<AllShaders, ShaderGroups> Pipeline;
typedef DescriptorSetLayoutVulkanConst<3, {{
    {
        .binding = 0,
        .descriptorType = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR,
        .descriptorCount = 1,
        .stageFlags = VK_SHADER_STAGE_RAYGEN_BIT_KHR,
    },
    {
        .binding = 1,
        .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
        .descriptorCount = 1,
        .stageFlags = VK_SHADER_STAGE_RAYGEN_BIT_KHR,
    },
    {
        .binding = 2,
        .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .descriptorCount = 1,
        .stageFlags = VK_SHADER_STAGE_RAYGEN_BIT_KHR,
    }
}}> descriptorSetLayout;

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

    Raygenspv::CreateShader();
    Closesthitspv::CreateShader();
    Misspv::CreateShader();
    ShaderBindingTableVulkanConst<AllShaders>::Init();

    descriptorSetLayout::Init();
    std::array<VkDescriptorSetLayout, 1> layouts {{descriptorSetLayout::layout}};

    DescriptorPool pool;
    pool.sets.resize(1);
    pool.BuildPool(DescriptorPool::GetPoolSizes<descriptorSetLayout>(), layouts);

    Pipeline::Init(cmd, layouts);

    Mesh triangleMesh;
    std::array<Vector<float, 3, 3>, 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 renderer = {
        .instance = {
            .instanceCustomIndex                    = 0,
            .mask                                   = 0xFF,
            .instanceShaderBindingTableRecordOffset = 0,
            .flags                                  = VK_GEOMETRY_INSTANCE_FORCE_OPAQUE_BIT_KHR,
            .accelerationStructureReference         = triangleMesh.blasAddr
        }
    };

    RenderingElement3DVulkan::elements.emplace_back(&renderer);

    MatrixRowMajor<float, 4, 3, 1> transform = MatrixRowMajor<float, 4, 3, 1>::Identity();
    std::memcpy(renderer.instance.transform.matrix, transform.m, sizeof(transform.m));
    RenderingElement3DVulkan::tlases.resize(1);
    RenderingElement3DVulkan::BuildTLAS(cmd, 0);

    VkDescriptorImageInfo imageInfo = {
        .imageView = window.imageViews[0],
        .imageLayout = VK_IMAGE_LAYOUT_GENERAL
    };

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

    VulkanBuffer<std::uint32_t, true, false, false> lightBuffer;
    lightBuffer.Create(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, 1);

    VkWriteDescriptorSet write[3] = {
        {
            .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,
        },
        {
            .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
            .dstSet = pool.sets[0],
            .dstBinding = 1,
            .dstArrayElement = 0,
            .descriptorCount = 1,
            .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
            .pImageInfo = &imageInfo
        },
        {
            .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
            .dstSet = pool.sets[0],
            .dstBinding = 2,
            .dstArrayElement = 0,
            .descriptorCount = 1,
            .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            .pBufferInfo = &lightBuffer.descriptor
        }
    };
    vkUpdateDescriptorSets(VulkanDevice::device, 3, write, 0, nullptr);
    

    window.SetPipelineRT<Pipeline>();
    window.descriptorsRt = pool.sets;

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