// RTStress — the standing many-instance wavefront RT benchmark. An
// N×N×N grid of a small cube mesh (one BLAS, many TLAS instances), shaded
// with primary + shadow rays through the wavefront pipeline. The grid edge
// `kGrid` is the instance-count knob: 8 → 512, 16 → 4096, 20 → 8000
// (LBVH_MAX = 16384). Frame time is printed to the console each second so
// fps-vs-instance-count can be read off without external tooling; the JS
// bridge additionally prints a GPU timestamp-query per-pass breakdown.
//
// WebGPU/DOM only — the wavefront tracer is the WebGPU software RT path.

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
int main() { return 0; }   // native path is hardware RT; out of scope here
#else

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

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

namespace {
    // Instance-count knob. instances = kGrid³. Bump to 16 (4096) or 20
    // (8000) to stress the TLAS; the LBVH build caps at 16384.
    constexpr int   kGrid    = 8;
    constexpr float kSpacing = 2.5f;
    constexpr float kHalf    = 0.5f;   // cube half-extent

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

int main() {
    const int instanceCount = kGrid * kGrid * kGrid;
    std::println("[RTStress] grid {}^3 = {} instances", kGrid, instanceCount);

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

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

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

    // One user binding: the camera storage buffer at @group(3).
    std::array<UICustomBinding, 1> bindings {{
        { .group = 3, .binding = 0, .kind = UICustomBindingKind::Buffer, .pushOffset = 0 },
    }};

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

    // ── Unit cube mesh (8 verts, 12 tris). ────────────────────────────
    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);

    // ── Camera buffer + handle array. ─────────────────────────────────
    WebGPUBuffer<CameraGPU, true> cameraBuf;
    cameraBuf.Create(1);
    static std::array<std::uint32_t, 1> userHandles { cameraBuf.handle };

    // ── Instance grid. Reserve so RenderingElement3D::Add pointers stay
    //    valid across vector growth. ─────────────────────────────────────
    static std::vector<RenderingElement3D> renderers;
    renderers.reserve(static_cast<std::size_t>(instanceCount));
    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);

    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 framing the grid from a corner. ───────────────────
    const float ext = float(kGrid - 1) * kSpacing;
    struct CamState {
        Vector<float, 3, 4> position;
        float yaw;
        float pitch;
    } cam {
        Vector<float, 3, 4>{ ext * 1.4f, ext * 1.0f, ext * 1.4f },
        0.0f, 0.0f,
    };
    {
        // Aim at the grid centre (origin).
        Vector<float, 3, 4> d { -cam.position.x, -cam.position.y, -cam.position.z };
        const float len = std::sqrt(d.x*d.x + d.y*d.y + d.z*d.z);
        cam.yaw   = std::atan2(d.z, d.x);
        cam.pitch = std::asin(d.y / len);
    }

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

    const float kMoveSpeed = ext * 0.8f;
    const float kLookSens  = 0.05f;
    const float kDt        = 1.0f / 60.0f;

    static int   frames = 0;
    static double tAccum = 0.0;
    EventListener<void> camTick(&window.onBeforeUpdate, [&]() {
        inputMap.Tick();
        cam.yaw   += lookAct.vector2.x * kLookSens;
        cam.pitch -= lookAct.vector2.y * kLookSens;
        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;
        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;
        cameraBuf.FlushDevice();

        if (++frames >= 60) {
            std::println("[RTStress] {} instances @ ~{} frames since last report", instanceCount, frames);
            frames = 0;
        }
    });

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