//SPDX-License-Identifier: LGPL-3.0-only
//SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

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

// DOM-mode TLAS upkeep. BuildTLAS is split in two phases so a physics
// compute pass can run between them:
//   - BuildTLASUpload mirrors the CPU-side RTInstance array into the
//     host-visible instance buffer (with partial-write semantics that
//     preserve the transform bytes for elements flagged
//     transformOwnedByGpu, see notes in the body) and uploads the
//     metadata buffer.
//   - BuildTLASBuild 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.
// The combined BuildTLAS calls both back-to-back; callers that want to
// interleave a physics tlas-transform compute pass (which writes the
// transform bytes BuildTLASUpload leaves intact) call Upload + their
// compute pass + Build manually.

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::BuildTLASUpload(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;
    }

    constexpr std::uint32_t kNPadded   = 65536u;     // size for instance / metadata mirrors
    constexpr std::uint32_t kLbvhMax   = 16384u;     // matches N_PADDED in lbvhBuildWgsl
    constexpr std::uint32_t kNodeCount = 2u * kNPadded - 1u;

    // ALL TLAS-side GPU buffers get allocated ONCE and never resized.
    // The LBVH-build shader takes the real instance count via a uniform
    // (lbvhPc.nReal) instead of arrayLength(&entries), so the
    // tlas.buffer / entryOrder / mortonCodes don't need to grow when
    // the application's element count changes.
    //
    // Why this matters: an earlier version resized these per-frame on
    // primitiveCount change. The destroy+recreate cycle on the GPU
    // buffer caused subtle mid-game flicker as soon as any element was
    // added (e.g. firing a projectile) — fort braces would appear to
    // briefly vanish in patterns deterministic on the projectile's
    // angle. Suspected driver-level memory recycling without proper
    // zero-init; the fixed-size allocation sidesteps it entirely.
    if (tlas.instanceBuffer.handle == 0) {
        tlas.instanceBuffer.Resize(kNPadded);
        tlas.metadataBuffer.Resize(kNPadded);
        tlas.bvhNodes.Resize(kNodeCount * 32u);
        tlas.sortTempA.Resize(kNPadded * 4u);
        tlas.sortTempB.Resize(kNPadded * 4u);
        tlas.tlasBins.Resize(64 * 32);
        // TLAS-entry / order / morton-code buffers: sized for the LBVH
        // cap (16384). lbvhBuildMain iterates `lbvhPc.nReal` real
        // entries; the remainder stays zero / sentinel. Keep these
        // stable across element-count changes so the renderer's bind
        // group references the same buffer handle every frame.
        tlas.buffer.Resize(kLbvhMax * 144u);
        tlas.entryOrder.Resize(kLbvhMax * 4u);
        tlas.mortonCodes.Resize(kLbvhMax * 4u);
    }

    // NB: tlas.buffer / entryOrder / mortonCodes get resized in
    // BuildTLASBuild, NOT here. Resize destroys + recreates the GPU
    // resource (and the JS-side handle); the rayQuery dispatches that
    // run between BuildTLASUpload and BuildTLASBuild (projectile-collide,
    // splash, builder-pick) still hold the previous frame's TLAS in
    // rtState.current{Tlas,EntryOrder,Bvh}. If we resized here, those
    // handles would point at destroyed buffers and the dispatches would
    // log "no TLAS built yet" every frame the element count changed
    // (e.g. every projectile fire). Resizing inside BuildTLASBuild,
    // immediately before wgpuBuildTLAS publishes the new handles, keeps
    // the JS-side current* refs in sync with the GPU resources.

    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) {
            // GPU owns the transform: copy only the host-authored
            // metadata fields and leave the host mirror's transform bytes
            // alone (the application doesn't maintain them). The upload
            // below pushes the whole struct, including these transform
            // bytes, but the physics-tlas-transform compute pass — which
            // runs after this Upload and before Build in the supported
            // Upload -> compute -> Build flow — overwrites the transform
            // on the GPU before the TLAS build reads it, so the host-side
            // value is moot. (Nothing reads instanceBuffer between Upload
            // and that compute pass.)
            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;
    }

    // Upload the active instance range in a single contiguous writeBuffer.
    //
    // An earlier version split this into per-run arms: contiguous runs of
    // CPU-driven elements (transformOwnedByGpu=false) were pushed as one
    // block, but GPU-driven runs pushed only the trailing 16 metadata
    // bytes of each element — one FlushDeviceRange per element — to avoid
    // clobbering the GPU-written transform. The metadata bytes are strided
    // (one 16-byte chunk per 64-byte slot), so a GPU-driven run couldn't
    // batch; each element paid a separate WebGPU validation / encode /
    // JS-boundary cost.
    //
    // Pushing the whole struct for GPU-driven elements too is harmless:
    // the only supported way to drive a transform from the GPU is the
    // manual Upload -> physics compute pass -> Build sequence (see the
    // file header), and that compute pass runs after this upload and
    // rewrites the transform on the GPU before the TLAS build reads it.
    // So the stale transform bytes this push carries for those slots are
    // overwritten before they matter, and the whole run uploads in one
    // call. (The combined BuildTLAS path runs no compute pass and is, as
    // documented, only valid for CPU-driven transforms.)
    constexpr std::uint32_t kInstSize = sizeof(RTInstance);   // 64
    tlas.instanceBuffer.FlushDeviceRange(0, 0, primitiveCount * kInstSize);

    // Only primitiveCount slots are live; the buffer is padded to kNPadded
    // (65536). Flushing the whole thing wgpuWriteBuffers all 256 KB through
    // the WASM->JS staging path every frame — ~100-250x waste for a
    // few-hundred-instance scene. Range-flush just the live entries.
    tlas.metadataBuffer.FlushDeviceRange(0, 0, primitiveCount * sizeof(std::uint32_t));
}

void RenderingElement3D::BuildTLASBuild(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) {
        // Upload already cleared builtInstanceCount; nothing to dispatch.
        return;
    }

    // No per-count Resize. tlas.buffer / entryOrder / mortonCodes were
    // allocated at kLbvhMax in BuildTLASUpload's first call and stay
    // that size. The LBVH shader reads the real count from a uniform
    // (lbvhPc.nReal) wgpuBuildTLAS writes each call.

    WebGPU::wgpuBuildTLAS(tlas.instanceBuffer.handle,
                          static_cast<std::int32_t>(primitiveCount),
                          tlas.buffer.handle,
                          tlas.entryOrder.handle,
                          tlas.mortonCodes.handle,
                          tlas.tlasBins.handle,
                          tlas.bvhNodes.handle,
                          tlas.sortTempA.handle,
                          tlas.sortTempB.handle);

    tlas.builtInstanceCount = primitiveCount;
}

void RenderingElement3D::BuildTLAS(WebGPUCommandEncoderRef cmd, std::uint32_t index, RTBuildPreference) {
    // `preference` is ignored — the software LBVH TLAS build has no
    // FAST_TRACE/FAST_BUILD equivalent. See the interface comment.
    BuildTLASUpload(cmd, index);
    BuildTLASBuild(cmd, index);
}