// HDRBloom — cross-backend-shaped HDR post-process on the WebGPU/DOM
// backend, exercising the primitives added for issue #27:
//
//   1. rgba16float runtime textures (StorageImage2D) with STORAGE + SAMPLED
//      usage — the scene + bloom-mip targets.
//   2. A write-only StorageTexture binding kind (UICustomBindingKind::
//      StorageTexture) so compute passes write those float targets, plus
//      float-texture sampling via SampledTexture.
//   3. PipelineRTWebGPU's RESOLVE writing linear radiance into a user
//      rgba16float texture (hdrOutputFormat) instead of the canvas.
//
// Pipeline shape (the same one a Vulkan bloom would use):
//   RT  → linear rgba16float scene
//   threshold (compute, sample float → write float)
//   blur      (compute, sample float → write float)
//   composite (UI custom shader: scene + bloom → tonemap+gamma → canvas)
//
// The threshold/blur passes run from onBeforeUpdate so each lands on its own
// queue submit — WebGPU's per-submit ordering gives the storage-write →
// sampled-read barrier the chain needs (there is no barrier between
// dispatches within one compute pass). The composite runs in-frame as a UI
// custom shader so it owns the canvas ping-pong. Bloom is therefore one
// frame behind the sharp scene, which is imperceptible for this static view.
//
// WebGPU/DOM only — the wavefront tracer is the WebGPU software RT path.

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
int main() { return 0; }   // native bloom is wireable today (see issue gap 4)
#else

#include <cstddef>  // offsetof

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

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

namespace {
    constexpr int   kGrid    = 3;
    constexpr float kSpacing = 2.5f;
    constexpr float kHalf    = 0.5f;

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

    struct Dim { std::uint32_t w, h, _0, _1; };

    // Composite push: standard header + the three heap slots the UI custom
    // shader's group(2) bindings resolve through.
    struct CompositePush {
        UIDispatchHeader hdr;
        std::uint32_t    hdrSlot;
        std::uint32_t    bloomSlot;
        std::uint32_t    sampSlot;
        std::uint32_t    _pad;
    };
}

int main() {
    Device::Initialize();
    static Window window(1280, 720, "HDRBloom");
    auto cmd = window.StartInit();

    DescriptorHeapWebGPU heap;
    heap.Initialize(/*images*/ 4, /*buffers*/ 4, /*samplers*/ 2);
    window.descriptorHeap = &heap;

    const std::uint16_t W = static_cast<std::uint16_t>(window.width);
    const std::uint16_t H = static_cast<std::uint16_t>(window.height);

    // ── RT pipeline: HDR output (RESOLVE → rgba16float) ────────────────
    std::array<WebGPUShader, 3> 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),
    }};
    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 } }};

    std::array<UICustomBinding, 1> rtBindings {{
        { .group = 3, .binding = 0, .kind = UICustomBindingKind::Buffer, .pushOffset = 0 },
    }};

    PipelineRTWebGPU pipeline;
    pipeline.Init(cmd, raygenGroups, missGroups, hitGroups, sbt, rtBindings,
                  WebGPUTexelFormat::RGBA16Float);   // ← RESOLVE writes HDR

    // ── Unit cube mesh. ────────────────────────────────────────────────
    static std::array<Vector<float, 3, 3>, 8> verts {{
        {-kHalf, -kHalf, -kHalf}, { kHalf, -kHalf, -kHalf},
        { kHalf,  kHalf, -kHalf}, {-kHalf,  kHalf, -kHalf},
        {-kHalf, -kHalf,  kHalf}, { kHalf, -kHalf,  kHalf},
        { kHalf,  kHalf,  kHalf}, {-kHalf,  kHalf,  kHalf},
    }};
    static std::array<std::uint32_t, 36> indices {{
        0,1,2, 0,2,3,   5,4,7, 5,7,6,   4,0,3, 4,3,7,
        1,5,6, 1,6,2,   4,5,1, 4,1,0,   3,2,6, 3,6,7,
    }};
    static Mesh cube;
    cube.Build(verts, indices, cmd);

    WebGPUBuffer<CameraGPU, true> cameraBuf;
    cameraBuf.Create(1);
    static std::array<std::uint32_t, 1> rtHandles { cameraBuf.handle };

    static std::vector<RenderingElement3D> renderers;
    renderers.reserve(static_cast<std::size_t>(kGrid * kGrid * kGrid));
    const float origin0 = -0.5f * static_cast<float>(kGrid - 1) * kSpacing;
    for (int x = 0; x < kGrid; ++x)
    for (int y = 0; y < kGrid; ++y)
    for (int z = 0; z < kGrid; ++z) {
        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] = origin0 + float(x) * kSpacing;
        tx[1][0] = 0; tx[1][1] = 1; tx[1][2] = 0; tx[1][3] = origin0 + float(y) * kSpacing;
        tx[2][0] = 0; tx[2][1] = 0; tx[2][2] = 1; tx[2][3] = origin0 + float(z) * kSpacing;
        r.instance.instanceCustomIndex                    = static_cast<std::uint32_t>(renderers.size() - 1);
        r.instance.mask                                   = 0xFF;
        r.instance.instanceShaderBindingTableRecordOffset = 0;
        r.instance.flags                                  = kRTGeometryInstanceForceOpaque;
        r.instance.accelerationStructureReference         = cube.blasAddr;
        RenderingElement3D::Add(&r);
    }
    RenderingElement3D::BuildTLAS(cmd, 0);

    // ── HDR scene + bloom mip targets (rgba16float). ───────────────────
    StorageImage2D hdrScene;  hdrScene.Create(W, H);
    StorageImage2D bloomA;    bloomA.Create(W, H);
    StorageImage2D bloomB;    bloomB.Create(W, H);

    // Heap slots for the composite's sampled inputs + sampler.
    ImageSlot   hdrSlot   = hdrScene.AllocateSlot(heap);
    ImageSlot   bloomSlot = bloomB.AllocateSlot(heap);
    SamplerSlot sampSlot  = AllocateLinearClampSampler(heap);

    // ── Threshold + blur compute passes. ───────────────────────────────
    std::array<UICustomBinding, 2> bloomBindings {{
        { .group = 1, .binding = 0, .kind = UICustomBindingKind::SampledTexture, .pushOffset = 0 },
        { .group = 1, .binding = 1, .kind = UICustomBindingKind::StorageTexture,
          .format = static_cast<std::uint8_t>(WebGPUTexelFormat::RGBA16Float), .pushOffset = 0 },
    }};
    PlainComputeShader threshold;
    threshold.Load(fs::path("threshold.comp.wgsl"), sizeof(Dim), bloomBindings);
    PlainComputeShader blur;
    blur.Load(fs::path("blur.comp.wgsl"), sizeof(Dim), bloomBindings);

    std::array<std::uint32_t, 2> thresholdHandles { hdrScene.handle, bloomA.handle };
    std::array<std::uint32_t, 2> blurHandles      { bloomA.handle,   bloomB.handle };

    // ── Composite UI custom shader. ────────────────────────────────────
    UIRenderer ui;
    ui.Initialize(window, heap, cmd);

    UICustomBinding compBindings[] = {
        { .group = 2, .binding = 0, .kind = UICustomBindingKind::SampledTexture,
          .pushOffset = static_cast<std::uint32_t>(offsetof(CompositePush, hdrSlot)) },
        { .group = 2, .binding = 1, .kind = UICustomBindingKind::SampledTexture,
          .pushOffset = static_cast<std::uint32_t>(offsetof(CompositePush, bloomSlot)) },
        { .group = 2, .binding = 2, .kind = UICustomBindingKind::Sampler,
          .pushOffset = static_cast<std::uint32_t>(offsetof(CompositePush, sampSlot)) },
    };
    WebGPUComputeShader composite;
    composite.Load(fs::path("composite.comp.wgsl"), compBindings);

    window.FinishInit();

    // ── Passes: RT first (writes HDR scene), then UI (composite→canvas).
    RTPass rtPass(&pipeline);
    rtPass.handlesPtr   = rtHandles.data();
    rtPass.handlesCount = static_cast<std::uint32_t>(rtHandles.size());
    rtPass.maxDepth     = 1;                 // primary rays only
    rtPass.outTexHandle = hdrScene.handle;   // ← RESOLVE target
    window.passes.push_back(&rtPass);
    window.passes.push_back(&ui);

    // ── Static camera framing the grid from a front corner. ────────────
    const float ext = float(kGrid - 1) * kSpacing;
    Vector<float, 3, 4> camPos { ext * 1.1f, ext * 0.8f, ext * 1.6f + 3.0f };
    Vector<float, 3, 4> d { -camPos.x, -camPos.y, -camPos.z };
    const float dl = std::sqrt(d.x*d.x + d.y*d.y + d.z*d.z);
    Vector<float, 3, 4> forward { d.x/dl, d.y/dl, d.z/dl };
    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 rl = std::sqrt(right.x*right.x + right.y*right.y + right.z*right.z);
    right.x /= rl; right.y /= rl; right.z /= rl;
    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 };
    {
        CameraGPU& g = cameraBuf.value[0];
        g.origin[0]=camPos.x; g.origin[1]=camPos.y; g.origin[2]=camPos.z; g.pad0=0;
        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(60.0f * 3.14159265f / 360.0f);
        g.pad1    = 0;
        cameraBuf.FlushDevice();
    }

    // ── Bloom prepass: threshold + blur, each its own submit. ──────────
    EventListener<void> bloomTick(&window.onBeforeUpdate, [&]() {
        Dim dim { static_cast<std::uint32_t>(window.width),
                  static_cast<std::uint32_t>(window.height), 0, 0 };
        const std::uint32_t gx = (window.width  + 7u) / 8u;
        const std::uint32_t gy = (window.height + 7u) / 8u;
        threshold.Dispatch(&dim, sizeof(dim), thresholdHandles, gx, gy, 1);
        blur.Dispatch(&dim, sizeof(dim), blurHandles, gx, gy, 1);
    });

    // ── Composite: scene + bloom → tonemap+gamma → canvas. ─────────────
    EventListener<UIBuildArgs> composeSub(&ui.onBuild, [&](UIBuildArgs a) {
        CompositePush pc { ui.FillHeader(0, 0), 0, 0, 0, 0 };
        pc.hdrSlot   = static_cast<std::uint32_t>(static_cast<std::uint16_t>(hdrSlot));
        pc.bloomSlot = static_cast<std::uint32_t>(static_cast<std::uint16_t>(bloomSlot));
        pc.sampSlot  = static_cast<std::uint32_t>(static_cast<std::uint16_t>(sampSlot));
        const std::uint32_t gx = (window.width  + 7u) / 8u;
        const std::uint32_t gy = (window.height + 7u) / 8u;
        ui.Dispatch(a.cmd, composite, &pc, sizeof(pc), gx, gy, 1);
    });

    std::println("[HDRBloom] RT→rgba16float→threshold→blur→composite running");

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