#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
#include "vulkan/vulkan.h"
#endif
#include <cassert>

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

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
int main() {
    Device::Initialize();
    Window window(1280, 720, "HelloVulkan");
    VkCommandBuffer cmd = window.StartInit();
    DescriptorHeapVulkan descriptorHeap;
    descriptorHeap.Initialize(1,1,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, 3> shaders{{
        {"raygen.spv", "main", VK_SHADER_STAGE_RAYGEN_BIT_KHR, &specilizationInfo},
        {"miss.spv", "main", VK_SHADER_STAGE_MISS_BIT_KHR, nullptr},
        {"closesthit.spv", "main", VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR, nullptr}
    }};

    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, 1> missGroups {{
        {
            .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
            .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR,
            .generalShader = 1,
            .closestHitShader = VK_SHADER_UNUSED_KHR,
            .anyHitShader = VK_SHADER_UNUSED_KHR,
            .intersectionShader = VK_SHADER_UNUSED_KHR,
        },
    }};
    std::array<VkRayTracingShaderGroupCreateInfoKHR, 1> hitGroups {{
        {
            .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
            .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR,
            .generalShader = VK_SHADER_UNUSED_KHR,
            .closestHitShader = 2,
            .anyHitShader = VK_SHADER_UNUSED_KHR,
            .intersectionShader = VK_SHADER_UNUSED_KHR,
        },
    }};

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

    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);

    RenderingElement3D renderer = {
        .instance = {
            .instanceCustomIndex                    = 0,
            .mask                                   = 0xFF,
            .instanceShaderBindingTableRecordOffset = 0,
            .flags                                  = VK_GEOMETRY_INSTANCE_FORCE_OPAQUE_BIT_KHR,
            .accelerationStructureReference         = triangleMesh.blasAddr
        }
    };

    RenderingElement3D::elements.emplace_back(&renderer);

    MatrixRowMajor<float, 4, 3, 1> transform = MatrixRowMajor<float, 4, 3, 1>::Identity();
    transform.Store(reinterpret_cast<float*>(renderer.instance.transform.matrix));
    RenderingElement3D::BuildTLAS(cmd, 0);
    RenderingElement3D::BuildTLAS(cmd, 1);
    RenderingElement3D::BuildTLAS(cmd, 2);

    window.FinishInit();

    auto imgSlots = descriptorHeap.AllocateImageSlots(1);
    auto bufSlots = descriptorHeap.AllocateBufferSlots(1);

    VkDeviceAddressRangeKHR tlasRange0 = {
        .address = RenderingElement3D::tlases[0].address,
    };

    VkDeviceAddressRangeKHR tlasRange1 = {
        .address = RenderingElement3D::tlases[1].address,
    };

    VkDeviceAddressRangeKHR tlasRange2 = {
        .address = RenderingElement3D::tlases[2].address,
    };

    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
    };

    VkResourceDescriptorInfoEXT resources[6] = {
        {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR,
            .data = { .pAddressRange =  &tlasRange0}
        },
        {
            .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_ACCELERATION_STRUCTURE_KHR,
            .data = { .pAddressRange =  &tlasRange1}
        },
        {
            .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_ACCELERATION_STRUCTURE_KHR,
            .data = { .pAddressRange =  &tlasRange2}
        },
        {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
            .data = { .pImage = &imageInfo2 }
        },
    };
    
    VkHostAddressRangeEXT destinations[6] = {
        {
            .address = descriptorHeap.resourceHeap[0].value + descriptorHeap.BufferByteOffset(bufSlots.firstElement),
            .size = Device::descriptorHeapProperties.bufferDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[0].value + descriptorHeap.ImageByteOffset(imgSlots.firstElement),
            .size = Device::descriptorHeapProperties.imageDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[1].value + descriptorHeap.BufferByteOffset(bufSlots.firstElement),
            .size = Device::descriptorHeapProperties.bufferDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[1].value + descriptorHeap.ImageByteOffset(imgSlots.firstElement),
            .size = Device::descriptorHeapProperties.imageDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[2].value + descriptorHeap.BufferByteOffset(bufSlots.firstElement),
            .size = Device::descriptorHeapProperties.bufferDescriptorSize
        },
        {
            .address = descriptorHeap.resourceHeap[2].value + descriptorHeap.ImageByteOffset(imgSlots.firstElement),
            .size = Device::descriptorHeapProperties.imageDescriptorSize
        },
    };

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

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

    window.descriptorHeap = &descriptorHeap;
    RTPass rtPass(&pipeline);
    window.passes.push_back(&rtPass);

    // NOTE: reading the acceleration structure through VK_EXT_descriptor_heap
    // aborts with VK_ERROR_DEVICE_LOST on NVIDIA 610.43.02 (a driver fault —
    // see #7). The engine transparently works around it: on the NVIDIA driver
    // VulkanShader rewrites the heap AS read into a TLAS-device-address +
    // OpConvertUToAccelerationStructureKHR path and RTPass feeds the address in
    // as push data. Nothing here (or in raygen.glsl) changes. See README.md
    // ("Native status") and interfaces/Crafter.Graphics-ShaderVulkan.cppm.
    window.Render();
    window.StartSync();
}
#else
// DOM-mode port. Same scene (one triangle), software-emulated raytracing
// via compute. Shaders are read from .wgsl files shipped as static
// assets (see project.cpp). Renders barycentric colors via the
// hit/miss/raygen mega-switch in PipelineRTWebGPU.
int main() {
    Device::Initialize();
    static Window window(1280, 720, "HelloVulkan");
    auto cmd = window.StartInit();

    DescriptorHeapWebGPU heap;
    heap.Initialize(/*images*/ 4, /*buffers*/ 4, /*samplers*/ 2);

    std::array<WebGPUShader, 3> shaders {{
        WebGPUShader(std::filesystem::path("raygen.wgsl"),     "raygen_main",     WebGPURTStage::Raygen),
        WebGPUShader(std::filesystem::path("miss.wgsl"),       "miss_main",       WebGPURTStage::Miss),
        WebGPUShader(std::filesystem::path("closesthit.wgsl"), "closesthit_main", WebGPURTStage::ClosestHit),
    }};
    ShaderBindingTableWebGPU sbt;
    sbt.Init(shaders);

    std::array<RTShaderGroup, 1> raygenGroups {{
        { .type = RTShaderGroupType::General,           .generalShader = 0 },
    }};
    std::array<RTShaderGroup, 1> missGroups {{
        { .type = RTShaderGroupType::General,           .generalShader = 1 },
    }};
    std::array<RTShaderGroup, 1> hitGroups {{
        { .type = RTShaderGroupType::TrianglesHitGroup, .closestHitShader = 2 },
    }};

    PipelineRTWebGPU pipeline;
    pipeline.Init(cmd, raygenGroups, missGroups, hitGroups, sbt);

    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);

    static RenderingElement3D renderer;
    renderer.instance.transform.matrix[0][0] = 1; renderer.instance.transform.matrix[0][1] = 0; renderer.instance.transform.matrix[0][2] = 0; renderer.instance.transform.matrix[0][3] = 0;
    renderer.instance.transform.matrix[1][0] = 0; renderer.instance.transform.matrix[1][1] = 1; renderer.instance.transform.matrix[1][2] = 0; renderer.instance.transform.matrix[1][3] = 0;
    renderer.instance.transform.matrix[2][0] = 0; renderer.instance.transform.matrix[2][1] = 0; renderer.instance.transform.matrix[2][2] = 1; renderer.instance.transform.matrix[2][3] = 0;
    renderer.instance.instanceCustomIndex                    = 0;
    renderer.instance.mask                                   = 0xFF;
    renderer.instance.instanceShaderBindingTableRecordOffset = 0;
    renderer.instance.flags                                  = kRTGeometryInstanceForceOpaque;
    renderer.instance.accelerationStructureReference         = triangleMesh.blasAddr;
    RenderingElement3D::Add(&renderer);
    RenderingElement3D::BuildTLAS(cmd, 0);

    window.descriptorHeap = &heap;
    window.FinishInit();

    RTPass rtPass(&pipeline);
    window.passes.push_back(&rtPass);

    window.Render();
    window.StartUpdate();
    window.StartSync();
}
#endif