/*
catcrafts.net
Copyright (C) 2026 Catcrafts

The source code of this website is made available for viewing purposes only.
No permission is granted to copy, modify, distribute, or create derivative works.

Ray-traced WebGPU demo embedded in the blog. The scene (a checkered floor
with five pillars lit by four coloured point lights, each pillar casting
four separable coloured shadows) and its four WGSL stages are adapted from
Crafter.Graphics' RTMultiShadow example — the wavefront software ray tracer
that is the WebGPU/DOM RT path. Differences from the example:

  * No input: the camera auto-orbits the scene, driven by a frame counter,
    so the demo just plays. The embedded canvas is offset from the viewport
    origin, which the mouse-coordinate bridge doesn't account for, so an
    interactive camera would be misaligned anyway.
  * The whole StartInit()/FinishInit() init lives in SetupDemo(); the render
    pass is only attached to the window while the demo is mounted into
    #webgpu-demo, so the rest of the site stays a plain DOM document.
*/

export module Catcrafts:Demo_impl;
import :Demo;
import Crafter.Graphics;
import Crafter.Math;
import Crafter.Event;
import std;

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

namespace {
    // Matches LIGHT_COUNT in shaders/closesthit.wgsl and drives the
    // per-pixel shadow-ray budget (one shadow ray per light).
    constexpr std::uint32_t kLightCount = 4;

    // Mirrors `struct Camera` in shaders/raygen.wgsl byte-for-byte.
    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);

    // Axis-aligned box: 8 corners between mn and mx.
    std::array<Vector<float, 3, 3>, 8> BoxVerts(float mnx, float mny, float mnz,
                                                float mxx, float mxy, float mxz) {
        return {{
            {mnx, mny, mnz}, {mxx, mny, mnz}, {mxx, mxy, mnz}, {mnx, mxy, mnz},
            {mnx, mny, mxz}, {mxx, mny, mxz}, {mxx, mxy, mxz}, {mnx, mxy, mxz},
        }};
    }
    // Mesh::Build takes mutable spans, so this can't be constexpr.
    std::array<std::uint32_t, 36> kBoxIndices {{
        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,
    }};

    // Cross-module demo state. Raw pointers into function-local statics in
    // SetupDemo() (valid for the program's lifetime); the site builds with
    // -fno-c++-static-destructors so none of it is torn down at exit.
    Window*           gWindow  = nullptr;
    RTPass*           gRtPass  = nullptr;
    bool              gActive  = false;
    std::uint32_t     gFrame   = 0;
}

namespace Catcrafts {

    void SetupDemo(Window& window) {
        gWindow = &window;

        auto cmd = window.StartInit();

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

        static 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),
        }};
        static ShaderBindingTableWebGPU sbt;
        sbt.Init(shaders);

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

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

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

        // Meshes: a large ground slab and a pillar (origin at its base).
        static auto groundVerts = BoxVerts(-30.0f, -1.0f, -30.0f, 30.0f, 0.0f, 30.0f);
        static auto pillarVerts = BoxVerts(-0.8f,   0.0f,  -0.8f,  0.8f, 6.0f,  0.8f);
        static Mesh ground, pillar;
        ground.Build(groundVerts, kBoxIndices, cmd);
        pillar.Build(pillarVerts, kBoxIndices, cmd);

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

        // Instances: ground (customIndex 0) + five pillars.
        struct Placement { float x, z; };
        static constexpr std::array<Placement, 5> kPillars {{
            { 0.0f, 0.0f }, { 5.0f, 5.0f }, { -5.0f, 5.0f }, { 5.0f, -5.0f }, { -5.0f, -5.0f },
        }};
        static std::vector<RenderingElement3D> renderers;
        renderers.reserve(1 + kPillars.size());
        auto addInstance = [&](std::uint64_t blasAddr, float x, float 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] = x;
            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] = z;
            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         = blasAddr;
            RenderingElement3D::Add(&r);
        };
        addInstance(ground.blasAddr, 0.0f, 0.0f);
        for (const auto& p : kPillars) addInstance(pillar.blasAddr, p.x, p.z);
        RenderingElement3D::BuildTLAS(cmd, 0);

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

        static RTPass rtPass(&pipeline);
        rtPass.handlesPtr   = userHandles.data();
        rtPass.handlesCount = static_cast<std::uint32_t>(userHandles.size());
        rtPass.maxDepth     = 2;            // primary + shadow
        rtPass.raysPerPixel = kLightCount;  // one shadow ray per light per pixel
        gRtPass = &rtPass;

        // Auto-orbit camera. Only does work (and touches the GPU buffer)
        // while the demo is mounted; otherwise the tick is a cheap no-op.
        // Capture nothing: cameraBuf is a static local (usable directly),
        // and the window is reached through the stable gWindow pointer —
        // capturing the reference parameter into this static lambda would
        // dangle once SetupDemo() returns.
        static EventListener<void> camTick(&window.onBeforeUpdate, []() {
            if (!gActive || gWindow == nullptr || gWindow->height == 0) return;

            const float t = static_cast<float>(gFrame) * 0.006f;
            ++gFrame;

            const float radius = 24.0f;
            const float px = std::cos(t) * radius;
            const float pz = std::sin(t) * radius;
            const float py = 12.0f;

            // forward = normalize(centre - eye), centre ≈ (0, 2, 0)
            float fx = -px, fy = 2.0f - py, fz = -pz;
            const float fl = std::sqrt(fx*fx + fy*fy + fz*fz);
            fx /= fl; fy /= fl; fz /= fl;

            // right = normalize(cross(forward, worldUp)), worldUp = (0,1,0)
            float rx = fy*0.0f - fz*1.0f;
            float ry = fz*0.0f - fx*0.0f;
            float rz = fx*1.0f - fy*0.0f;
            const float rl = std::sqrt(rx*rx + ry*ry + rz*rz);
            rx /= rl; ry /= rl; rz /= rl;

            // up = cross(right, forward)
            const float ux = ry*fz - rz*fy;
            const float uy = rz*fx - rx*fz;
            const float uz = rx*fy - ry*fx;

            CameraGPU& g = cameraBuf.value[0];
            g.origin[0]=px; g.origin[1]=py; g.origin[2]=pz; g.pad0=0;
            g.right[0]=rx;  g.right[1]=ry;  g.right[2]=rz;
            g.up[0]=ux;     g.up[1]=uy;     g.up[2]=uz;
            g.forward[0]=fx; g.forward[1]=fy; g.forward[2]=fz;
            g.aspect  = static_cast<float>(gWindow->width) / static_cast<float>(gWindow->height);
            g.tanHalf = std::tan(70.0f * 3.14159265f / 360.0f);
            g.pad1    = 0;
            cameraBuf.FlushDevice();
        });

        // Start as a plain DOM page: canvas detached + hidden, no RT pass.
        WebGPU::SetCanvasMount("");
    }

    void MountDemo() {
        if (!gWindow || !gRtPass) return;
        if (!gActive) {
            gWindow->passes.push_back(gRtPass);
            gActive = true;
        }
        WebGPU::SetCanvasMount(kDemoMountId);
    }

    void UnmountDemo() {
        if (!gWindow || !gRtPass) return;
        if (gActive) {
            auto& p = gWindow->passes;
            std::erase(p, gRtPass);
            gActive = false;
        }
        WebGPU::SetCanvasMount("");
    }
}