// Sponza on Vulkan + WebGPU. Same example source, two backends — picked
// by CRAFTER_GRAPHICS_WINDOW_DOM. Both paths:
//   1. Load a Sponza .cmesh (positions + indices, optional per-vertex
//      data region) and a single albedo .ctex from disk. The source
//      assets are fetched once by project.cpp (Crafter.Build::GitFetch)
//      from https://github.com/jimmiebergmann/Sponza and compressed
//      into the bin dir at build time — they don't live in this repo.
//   2. Build BLAS + TLAS via the existing Mesh / RenderingElement3D
//      flow. The on-disk format is identical between backends; only
//      the decompression path differs (VK_EXT_memory_decompression
//      on Vulkan, CPU GDeflate on WebGPU).
//   3. Upload the albedo as Image2D<RGBA8>, register it in the
//      backend descriptor heap, and run the RT pipeline. Closest-hit
//      shaders sample the texture at the hit's barycentric coords —
//      proof-of-binding rather than UV-correct shading. Per-vertex
//      UV interpolation is follow-up work (the attribs heap is in
//      place on WebGPU; the Vulkan side needs a sibling data buffer
//      exposed off Mesh).
//
// Sponza model: CC BY 3.0 — Frank Meinl (Crytek), packaged by Jimmie
// Bergmann and Morgan McGuire. https://casual-effects.com/data

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
#include "vulkan/vulkan.h"
#endif

import Crafter.Graphics;
import Crafter.Asset;
import Crafter.Math;
import Crafter.Event;
import std;

using namespace Crafter;
namespace fs = std::filesystem;

namespace {
    struct RGBA8 { std::uint8_t r, g, b, a; };

    void RequireAssets(const fs::path& mesh, const fs::path& tex) {
        const bool haveMesh = fs::exists(mesh);
        const bool haveTex  = fs::exists(tex);
        if (haveMesh && haveTex) return;
        std::println(std::cerr,
            "[Sponza] missing asset(s):\n"
            "  mesh:    {} {}\n"
            "  albedo:  {} {}\n"
            "The build should have populated these via cfg.assets +\n"
            "GitFetch (see examples/Sponza/project.cpp). If you ran\n"
            "the binary from outside its bin dir, cd into the bin dir\n"
            "first — asset paths are relative to cwd.",
            mesh.string(), haveMesh ? "OK" : "MISSING",
            tex.string(),  haveTex  ? "OK" : "MISSING");
        std::abort();
    }
}

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
int main() {
    // Native Vulkan path is single-material for now (see file header) —
    // pick up just the first per-material output the build emits. The
    // WebGPU branch below uses every mesh + a texture array.
    const fs::path meshPath = "mesh_0.cmesh";
    const fs::path texPath  = "tex_0.ctex";
    RequireAssets(meshPath, texPath);

    CompressedMeshAsset    loadedMesh = LoadCompressedMesh(meshPath);
    CompressedTextureAsset loadedTex  = LoadCompressedTexture(texPath);
    std::println("[Sponza] loaded {} verts, {} idx, {}x{} albedo",
                 loadedMesh.vertexCount, loadedMesh.indexCount,
                 loadedTex.sizeX, loadedTex.sizeY);

    Device::Initialize();
    Window window(1280, 720, "Sponza");
    VkCommandBuffer cmd = window.StartInit();

    DescriptorHeapVulkan descriptorHeap;
    descriptorHeap.Initialize(/*images*/ 2, /*buffers*/ 1, /*samplers*/ 0);

    // Two specialization constants: the TLAS slot offset (shared with
    // VulkanTriangle pattern) and the albedo slot index for closesthit.
    VkSpecializationMapEntry raygenEntry = { .constantID = 0, .offset = 0, .size = sizeof(std::uint16_t) };
    VkSpecializationInfo raygenSpec = {
        .mapEntryCount = 1, .pMapEntries = &raygenEntry,
        .dataSize = sizeof(std::uint16_t), .pData = &descriptorHeap.bufferStartElement,
    };

    // Allocate the albedo slot first so its index is known when we
    // compile closesthit.spv.
    auto imgSlots    = descriptorHeap.AllocateImageSlots(2);
    auto bufSlots    = descriptorHeap.AllocateBufferSlots(1);
    std::uint16_t albedoHeapSlot = static_cast<std::uint16_t>(imgSlots.firstElement + 1);

    VkSpecializationMapEntry hitEntry = { .constantID = 0, .offset = 0, .size = sizeof(std::uint16_t) };
    VkSpecializationInfo hitSpec = {
        .mapEntryCount = 1, .pMapEntries = &hitEntry,
        .dataSize = sizeof(std::uint16_t), .pData = &albedoHeapSlot,
    };

    std::array<VulkanShader, 3> shaders {{
        { "raygen.spv",     "main", VK_SHADER_STAGE_RAYGEN_BIT_KHR,      &raygenSpec },
        { "miss.spv",       "main", VK_SHADER_STAGE_MISS_BIT_KHR,        nullptr     },
        { "closesthit.spv", "main", VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR, &hitSpec    },
    }};
    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 sponzaMesh;
    sponzaMesh.Build(loadedMesh, cmd);

    Image2D<RGBA8> albedo;
    albedo.Create(loadedTex.sizeX, loadedTex.sizeY, /*mipLevels*/ 1, cmd,
                  VK_FORMAT_R8G8B8A8_UNORM,
                  VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
                  VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    albedo.Update(loadedTex, cmd, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    SamplerVulkan<RGBA8> sampler;

    static RenderingElement3D renderer;
    renderer.instance = {
        .transform                              = {},
        .instanceCustomIndex                    = 0,
        .mask                                   = 0xFF,
        .instanceShaderBindingTableRecordOffset = 0,
        .flags                                  = VK_GEOMETRY_INSTANCE_FORCE_OPAQUE_BIT_KHR,
        .accelerationStructureReference         = sponzaMesh.blasAddr,
    };
    MatrixRowMajor<float, 4, 3, 1>::Identity()
        .Store(reinterpret_cast<float*>(renderer.instance.transform.matrix));
    RenderingElement3D::elements.emplace_back(&renderer);
    RenderingElement3D::BuildTLAS(cmd, 0);
    RenderingElement3D::BuildTLAS(cmd, 1);
    RenderingElement3D::BuildTLAS(cmd, 2);

    window.FinishInit();

    // Write descriptors: TLAS at bufSlots[0], output image at imgSlots[0],
    // albedo (combined image+sampler) at imgSlots[1]. Per-frame replicated.
    VkDeviceAddressRangeKHR tlasRanges[Window::numFrames];
    VkImageDescriptorInfoEXT outImgInfos[Window::numFrames];
    VkDescriptorImageInfo albedoInfo {
        .sampler = sampler.textureSampler,
        .imageView = albedo.imageView,
        .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
    };
    for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
        tlasRanges[f] = { .address = RenderingElement3D::tlases[f].address };
        outImgInfos[f] = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_DESCRIPTOR_INFO_EXT,
            .pView = &window.imageViews[f],
            .layout = VK_IMAGE_LAYOUT_GENERAL,
        };
    }

    std::vector<VkResourceDescriptorInfoEXT> resources;
    std::vector<VkHostAddressRangeEXT>       destinations;
    resources.reserve(Window::numFrames * 3);
    destinations.reserve(Window::numFrames * 3);
    for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
        resources.push_back({
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR,
            .data = { .pAddressRange = &tlasRanges[f] },
        });
        destinations.push_back({
            .address = descriptorHeap.resourceHeap[f].value
                     + descriptorHeap.BufferByteOffset(bufSlots.firstElement),
            .size = Device::descriptorHeapProperties.bufferDescriptorSize,
        });
        resources.push_back({
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
            .data = { .pImage = &outImgInfos[f] },
        });
        destinations.push_back({
            .address = descriptorHeap.resourceHeap[f].value
                     + descriptorHeap.ImageByteOffset(imgSlots.firstElement),
            .size = Device::descriptorHeapProperties.imageDescriptorSize,
        });
        resources.push_back({
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
            .data = { .pCombinedImageSampler = &albedoInfo },
        });
        destinations.push_back({
            .address = descriptorHeap.resourceHeap[f].value
                     + descriptorHeap.ImageByteOffset(albedoHeapSlot),
            .size = Device::descriptorHeapProperties.imageDescriptorSize,
        });
    }
    Device::vkWriteResourceDescriptorsEXT(Device::device,
        static_cast<std::uint32_t>(resources.size()),
        resources.data(), destinations.data());
    for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
        descriptorHeap.resourceHeap[f].FlushDevice();
    }

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

    window.Render();
    window.StartSync();
    return 0;
}
#else
int main() {
    // ── Read scene manifest (produced by project.cpp's ImportSponzaBundle).
    //
    //   line 1: albedoCount
    //   line 2: meshCount
    //   line 3..: per-mesh albedoIdx (-1 means "no albedo")
    const fs::path manifestPath = "scene.txt";
    if (!fs::exists(manifestPath)) {
        std::println(std::cerr,
            "[Sponza] missing scene.txt — the build should have produced "
            "it (see examples/Sponza/project.cpp). If you ran the binary "
            "from outside its bin dir, cd in first.");
        std::abort();
    }
    std::ifstream manifest(manifestPath);
    std::uint32_t albedoCount = 0, meshCount = 0;
    manifest >> albedoCount >> meshCount;
    std::vector<std::int32_t> meshAlbedo(meshCount);
    for (std::uint32_t i = 0; i < meshCount; ++i) manifest >> meshAlbedo[i];
    std::println("[Sponza] scene: {} albedos, {} meshes", albedoCount, meshCount);

    Device::Initialize();
    static Window window(1280, 720, "Sponza");
    auto cmd = window.StartInit();

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

    std::array<WebGPUShader, 4> shaders {{
        WebGPUShader(fs::path("raygen.wgsl"),     "raygen_main",     WebGPURTStage::Raygen),
        WebGPUShader(fs::path("miss.wgsl"),       "miss_main",       WebGPURTStage::Miss),
        WebGPUShader(fs::path("closesthit.wgsl"), "closesthit_main", WebGPURTStage::ClosestHit),
        WebGPUShader(fs::path("resolve.wgsl"),    "resolve_main",    WebGPURTStage::Resolve),
    }};
    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 },
    }};

    // Three user bindings at @group(3) (the wavefront pipeline reserves
    // groups 0..2 for WfParams / data heaps / indirect args):
    //   binding 0 — albedo texture_2d_array (one layer per material)
    //   binding 1 — sampler (linear clamp)
    //   binding 2 — Camera storage buffer (host-driven, updated per frame)
    std::array<UICustomBinding, 3> bindings {{
        { .group = 3, .binding = 0, .kind = UICustomBindingKind::SampledTextureArray, ._pad = 0, .pushOffset = 0 },
        { .group = 3, .binding = 1, .kind = UICustomBindingKind::Sampler,             ._pad = 0, .pushOffset = 0 },
        { .group = 3, .binding = 2, .kind = UICustomBindingKind::Buffer,              ._pad = 0, .pushOffset = 0 },
    }};

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

    // ── Albedo texture array — one rgba8unorm layer per material. ──────
    //
    // Probe layer 0 for the canonical layer dimensions; project.cpp
    // already resized every albedo to the same square so any tex_N.ctex
    // would do, layer 0 is just the first one we have.
    Image2DArray<RGBA8> albedoArray;
    {
        CompressedTextureAsset probe = LoadCompressedTexture("tex_0.ctex");
        albedoArray.Create(probe.sizeX, probe.sizeY, static_cast<std::uint16_t>(albedoCount));
        albedoArray.UpdateLayer(0, probe);
        for (std::uint32_t i = 1; i < albedoCount; ++i) {
            CompressedTextureAsset tex = LoadCompressedTexture(std::format("tex_{}.ctex", i));
            albedoArray.UpdateLayer(static_cast<std::uint16_t>(i), tex);
        }
    }
    auto albedoArraySlot = albedoArray.AllocateSlot(heap);
    SamplerSlot samplerSlot = AllocateLinearClampSampler(heap);

    // Camera storage buffer — host writes (origin, right, up, forward,
    // aspect, tanHalf) every frame from the input-driven free camera
    // below. Layout matches the WGSL Camera struct in raygen.wgsl
    // (vec3-aligned, std430). 64 bytes total.
    struct CameraGPU {
        float origin[3];   float pad0;
        float right[3];    float tanHalf;
        float up[3];       float aspect;
        float forward[3];  float pad1;
    };
    static_assert(sizeof(CameraGPU) == 64);
    WebGPUBuffer<CameraGPU, true> cameraBuf;
    cameraBuf.Create(1);

    // Handle array fed to RTPass — order matches the bindings declaration.
    static std::array<std::uint32_t, 3> userHandles {
        heap.imageTable  [albedoArraySlot.firstElement],
        heap.samplerTable[samplerSlot.firstElement],
        cameraBuf.handle,
    };

    // ── Meshes + scene instances ───────────────────────────────────────
    //
    // One Mesh + one RenderingElement3D per material group from
    // scene.txt. Meshes whose albedoIdx is -1 (the .obj's `usemtl` named
    // something without a map_Kd in .mtl) get dropped — they're rare in
    // Sponza and we'd have nothing to sample for them anyway.
    //
    // Vector capacity is reserved up-front: RenderingElement3D::Add
    // takes a pointer that's stored in the static elements[] vector, so
    // any later vector reallocation would dangle those pointers.
    static std::vector<Mesh> meshes;
    static std::vector<RenderingElement3D> renderers;
    meshes.reserve(meshCount);
    renderers.reserve(meshCount);

    for (std::uint32_t i = 0; i < meshCount; ++i) {
        if (meshAlbedo[i] < 0) continue;
        CompressedMeshAsset loaded = LoadCompressedMesh(std::format("mesh_{}.cmesh", i));
        meshes.emplace_back();
        meshes.back().Build(loaded, cmd);

        renderers.emplace_back();
        RenderingElement3D& r = renderers.back();
        auto& tx = r.instance.transform.matrix;
        tx[0][0] = 1; tx[0][1] = 0; tx[0][2] = 0; tx[0][3] = 0;
        tx[1][0] = 0; tx[1][1] = 1; tx[1][2] = 0; tx[1][3] = 0;
        tx[2][0] = 0; tx[2][1] = 0; tx[2][2] = 1; tx[2][3] = 0;
        // 24-bit instanceCustomIndex carries the albedo array layer that
        // closesthit.wgsl reads as `hit.customIndex`.
        r.instance.instanceCustomIndex                    = static_cast<std::uint32_t>(meshAlbedo[i]);
        r.instance.mask                                   = 0xFF;
        r.instance.instanceShaderBindingTableRecordOffset = 0;
        r.instance.flags                                  = kRTGeometryInstanceForceOpaque;
        r.instance.accelerationStructureReference         = meshes.back().blasAddr;
        RenderingElement3D::Add(&r);
    }
    RenderingElement3D::BuildTLAS(cmd, 0);

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

    RTPass rtPass(&pipeline);
    rtPass.handlesPtr   = userHandles.data();
    rtPass.handlesCount = static_cast<std::uint32_t>(userHandles.size());
    rtPass.maxDepth     = 2;   // primary + shadow
    window.passes.push_back(&rtPass);

    // ── Free camera: WASD + mouse-delta look ───────────────────────────
    //
    // Initial pose puts the camera near one end of the atrium at eye
    // height, looking +X down the long axis (bbox: X[-1921..1800],
    // Y[-126..1429], Z[-1183..1105]). The user can fine-tune from there.
    struct CamState {
        // 3/4 view from a corner aimed at the atrium centre.
        Vector<float, 3, 4> position{ -1400.0f, 700.0f, -600.0f };
        float yaw   = 0.405f;   // radians, around world +Y
        float pitch = -0.317f;  // radians, +pitch looks up
    } cam;

    Input::Map inputMap;
    Input::Action& moveAct = inputMap.AddAction("Move", Input::ActionType::Vector2);
    Input::Action& lookAct = inputMap.AddAction("Look", Input::ActionType::Vector2);
    moveAct.bindings = {
        Input::WASDBind{
            Key(CrafterKeys::W), Key(CrafterKeys::S),
            Key(CrafterKeys::A), Key(CrafterKeys::D),
        },
    };
    lookAct.bindings = {
        Input::MouseDeltaBind{ 1.0f },
    };
    inputMap.Attach(window);

    constexpr float kMoveSpeed = 1200.0f;  // Sponza units / second (room is ~3700 wide)
    constexpr float kLookSens  = 0.05f;   // radians per mouse pixel
    constexpr float kDt        = 1.0f / 60.0f;

    EventListener<void> camTick(&window.onBeforeUpdate, [&]() {
        inputMap.Tick();

        cam.yaw   += lookAct.vector2.x * kLookSens;
        cam.pitch -= lookAct.vector2.y * kLookSens;
        // Keep pitch just shy of straight up/down so the basis vectors
        // don't collapse (cross(forward, world_up) would go zero).
        cam.pitch = std::clamp(cam.pitch, -1.55f, 1.55f);

        const float cp = std::cos(cam.pitch), sp = std::sin(cam.pitch);
        const float cy = std::cos(cam.yaw),   sy = std::sin(cam.yaw);
        Vector<float, 3, 4> forward { cp * cy, sp, cp * sy };
        Vector<float, 3, 4> worldUp { 0.0f, 1.0f, 0.0f };
        Vector<float, 3, 4> right { forward.y * worldUp.z - forward.z * worldUp.y,
                                    forward.z * worldUp.x - forward.x * worldUp.z,
                                    forward.x * worldUp.y - forward.y * worldUp.x };
        const float rLen = std::sqrt(right.x*right.x + right.y*right.y + right.z*right.z);
        right.x /= rLen; right.y /= rLen; right.z /= rLen;
        Vector<float, 3, 4> up { right.y * forward.z - right.z * forward.y,
                                 right.z * forward.x - right.x * forward.z,
                                 right.x * forward.y - right.y * forward.x };

        const float dx = moveAct.vector2.x * kMoveSpeed * kDt;
        const float dy = moveAct.vector2.y * kMoveSpeed * kDt;
        cam.position.x += right.x * dx + forward.x * dy;
        cam.position.y += right.y * dx + forward.y * dy;
        cam.position.z += right.z * dx + forward.z * dy;

        CameraGPU& g  = cameraBuf.value[0];
        g.origin[0]   = cam.position.x; g.origin[1]   = cam.position.y; g.origin[2]   = cam.position.z; g.pad0 = 0.0f;
        g.right[0]    = right.x;        g.right[1]    = right.y;        g.right[2]    = right.z;
        g.up[0]       = up.x;           g.up[1]       = up.y;           g.up[2]       = up.z;
        g.forward[0]  = forward.x;      g.forward[1]  = forward.y;      g.forward[2]  = forward.z;
        g.aspect      = float(window.width) / float(window.height);
        g.tanHalf     = std::tan(70.0f * 3.14159265f / 360.0f);
        g.pad1        = 0.0f;
        cameraBuf.FlushDevice();
    });

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