/*
Crafter®.Graphics
Copyright (C) 2026 Catcrafts®
catcrafts.net
*/

// DOM-mode TLAS upkeep. BuildTLAS copies the per-element RTInstance into
// the host-visible instance buffer (skipping the transform for elements
// whose transform is GPU-owned), uploads it, then dispatches the JS-side
// TLAS-build compute pass — which consults the per-BLAS records published
// at Mesh::Build() time to produce world-space AABBs and inverse
// transforms in the format `traceRay` / `rayQuery` consume.

module;
module Crafter.Graphics:RenderingElement3D_implWebGPU;

import :RenderingElement3D;
import :Mesh;
import :WebGPU;
import :WebGPUBuffer;
import std;

using namespace Crafter;

std::vector<RenderingElement3D*> RenderingElement3D::elements;

void RenderingElement3D::Add(RenderingElement3D* e) {
    e->indexInElements = static_cast<std::uint32_t>(elements.size());
    elements.push_back(e);
}

void RenderingElement3D::Remove(RenderingElement3D* e) {
    std::uint32_t idx = e->indexInElements;
    if (idx == std::numeric_limits<std::uint32_t>::max()) return;
    std::uint32_t last = static_cast<std::uint32_t>(elements.size() - 1);
    if (idx != last) {
        elements[idx] = elements[last];
        elements[idx]->indexInElements = idx;
    }
    elements.pop_back();
    e->indexInElements = std::numeric_limits<std::uint32_t>::max();
}

void RenderingElement3D::BuildTLAS(WebGPUCommandEncoderRef /*cmd*/, std::uint32_t index) {
    auto& tlas = tlases[index];
    const std::uint32_t primitiveCount = static_cast<std::uint32_t>(elements.size());
    if (primitiveCount == 0) {
        tlas.builtInstanceCount = 0;
        return;
    }

    // (Re)allocate instance + metadata + output TLAS buffers if the count
    // changed. WebGPUBuffer::Resize destroys and recreates the GPU buffer;
    // bind-group caches keyed on the buffer handle are invalidated in the
    // JS bridge automatically.
    if (primitiveCount != tlas.builtInstanceCount) {
        tlas.instanceBuffer.Resize(primitiveCount);
        tlas.metadataBuffer.Resize(primitiveCount);
        // TLASEntry layout in WGSL is 144 bytes due to vec3 align/pad
        // rules. Must match the struct declared in the rtWgslTypes
        // block in additional/dom-webgpu.js.
        tlas.buffer.Resize(primitiveCount * 144);
    }

    for (std::uint32_t i = 0; i < primitiveCount; ++i) {
        auto& dst = tlas.instanceBuffer.value[i];
        const auto& src = elements[i]->instance;
        if (elements[i]->transformOwnedByGpu) {
            // Preserve whatever the GPU compute shader most recently
            // wrote into dst.transform. Update only the non-transform
            // fields.
            dst.instanceCustomIndex                    = src.instanceCustomIndex;
            dst.mask                                   = src.mask;
            dst.instanceShaderBindingTableRecordOffset = src.instanceShaderBindingTableRecordOffset;
            dst.flags                                  = src.flags;
            dst.accelerationStructureReference         = src.accelerationStructureReference;
        } else {
            dst = src;
        }
        tlas.metadataBuffer.value[i] = elements[i]->userMetadata;
    }

    tlas.instanceBuffer.FlushDevice();
    tlas.metadataBuffer.FlushDevice();

    WebGPU::wgpuBuildTLAS(tlas.instanceBuffer.handle,
                          static_cast<std::int32_t>(primitiveCount),
                          tlas.buffer.handle);

    tlas.builtInstanceCount = primitiveCount;
}