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

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License version 3.0 as published by the Free Software Foundation;

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
*/

module;

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
#include "vulkan/vulkan.h"

#endif // !CRAFTER_GRAPHICS_WINDOW_DOM
export module Crafter.Graphics:Mesh;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
import std;
import Crafter.Math;
import Crafter.Asset;
import :VulkanBuffer;

export namespace Crafter {
    // One procedural primitive's axis-aligned bounding box, in object
    // space. Layout-compatible with VkAabbPositionsKHR (24-byte stride),
    // so a span of these uploads directly as a
    // VK_GEOMETRY_TYPE_AABBS_KHR build input. Mirrors the DOM-side
    // definition below so portable user code compiles unchanged.
    struct RTAabb {
        float min[3];
        float max[3];
    };
    static_assert(sizeof(RTAabb) == sizeof(VkAabbPositionsKHR));

    // Build-time tuning for a BLAS, mapped onto the preference bits of
    // VkBuildAccelerationStructureFlagBitsKHR.
    enum class RTBuildPreference {
        // VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR — spend
        // more time building for faster traversal. The right default for
        // static geometry that is traced many times.
        FastTrace,
        // VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_BUILD_BIT_KHR — build
        // quickly at the cost of traversal speed. For geometry rebuilt
        // (or first-frame streamed) often enough that build time dominates.
        FastBuild,
    };

    // Per-build options for Mesh::Build / BuildProcedural.
    struct RTBuildOptions {
        RTBuildPreference preference = RTBuildPreference::FastTrace;
        // Set VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_UPDATE_BIT_KHR so the
        // BLAS can later be refit in place via Refit()/RefitProcedural().
        // Required for any subsequent refit; refitting a BLAS built without
        // it falls back to a full rebuild. PREFER_FAST_TRACE and
        // PREFER_FAST_BUILD are mutually exclusive, so `preference` selects
        // exactly one — allowUpdate is layered on top of whichever is set.
        bool allowUpdate = false;
    };

    class Mesh {
        public:
        VulkanBuffer<char, false> scratchBuffer;
        VulkanBuffer<char, false> blasBuffer;
        VulkanBuffer<Vector<float, 3, 3>, true> vertexBuffer;
        VulkanBuffer<std::uint32_t, true> indexBuffer;
        // AABB build input for the procedural path (BuildProcedural).
        // Lifetime contract matches vertexBuffer/indexBuffer: must stay
        // alive until the build submitted on `cmd` completes.
        VulkanBuffer<RTAabb, true> aabbBuffer;
        // Lives until the cmd buffer issued by the compressed Build path
        // completes execution. Kept as a member so the recorded
        // vkCmdDecompressMemoryEXT references valid memory until the queue
        // submit signals — caller must not re-Build or destroy the Mesh
        // before that submit's fence is signaled (same contract as the
        // existing uncompressed path).
        VulkanBuffer<std::byte, true> compressedStaging;
        VkAccelerationStructureGeometryTrianglesDataKHR blasData;
        VkAccelerationStructureGeometryKHR blas;
        VkAccelerationStructureKHR accelerationStructure = VK_NULL_HANDLE;
        VkDeviceAddress blasAddr;
        bool opaque;
        // ─── Build-option / refit state ──────────────────────────────────
        // Flags the BLAS was last (re)built with — carried so Refit can
        // re-issue an identical-flags UPDATE build (the FAST_TRACE/FAST_BUILD
        // preference plus, when opted in, ALLOW_UPDATE).
        VkBuildAccelerationStructureFlagsKHR buildFlags = 0;
        // Primitive count of the current build (triangles = indexCount/3,
        // procedural = aabb count). An in-place UPDATE must preserve it.
        std::uint32_t builtPrimitiveCount = 0;
        // Element count the vertex / aabb input buffer was sized for. Refit
        // re-uploads in place only when the new data matches this; a change
        // forces a full rebuild (topology changed).
        std::uint32_t builtInputCount = 0;
        // Whether ALLOW_UPDATE was set on the last build, i.e. whether an
        // in-place refit is possible at all.
        bool allowUpdate = false;
        void Build(std::span<Vector<float, 3, 3>> verticies, std::span<std::uint32_t> indicies, VkCommandBuffer cmd, RTBuildOptions options = {});
        // GPU path: decompresses vertex (region 0) and index (region 1) streams
        // from asset.blob into vertexBuffer / indexBuffer using
        // VK_EXT_memory_decompression. Falls back to CPU decode + the
        // uncompressed Build if Device::memoryDecompressionSupported is false.
        // Region 2 (data) is not consumed here — the caller decompresses it
        // into their own buffer if needed (or uses Compression::DecompressCPU).
        void Build(const ::Crafter::CompressedMeshAsset& asset, VkCommandBuffer cmd, RTBuildOptions options = {});
        // Build an AABB (procedural) BLAS from a list of object-space
        // boxes (VK_GEOMETRY_TYPE_AABBS_KHR). The hit group bound to
        // instances of this mesh must be a
        // VK_RAY_TRACING_SHADER_GROUP_TYPE_PROCEDURAL_HIT_GROUP_KHR group
        // carrying an intersection shader; that shader is invoked for
        // each box the ray enters and reports the surface hit via
        // reportIntersectionEXT. `opaque` maps to
        // VK_GEOMETRY_OPAQUE_BIT_KHR: pass false (the WebGPU-path
        // default) to let any-hit shaders run. Same scratch/lifetime
        // handling as the triangle Build; instances reference blasAddr
        // exactly like a triangle BLAS.
        void BuildProcedural(std::span<const RTAabb> aabbs, bool opaque, VkCommandBuffer cmd, RTBuildOptions options = {});

        // Zero-copy procedural build: take the AABB build input straight from
        // an existing device buffer (by device address + primitive count)
        // rather than memcpy-ing a host span through aabbBuffer. The intended
        // producer is a GPU compute pass that writes the boxes itself — e.g.
        // a GPU-resident particle system whose per-frame AABBs never touch the
        // CPU. The buffer must carry
        // VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR
        // and VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, its boxes laid out as
        // VkAabbPositionsKHR (`stride`, default sizeof(RTAabb) == 24), and the
        // caller must pipeline-barrier the producing writes (e.g. a compute
        // dispatch) before the build on `cmd` reads them. The buffer's
        // lifetime is the caller's responsibility — it is never copied, so it
        // must outlive the build submitted on `cmd`. Everything else (opaque
        // bit, hit-group contract, instance wiring) matches the host-span
        // BuildProcedural above.
        void BuildProcedural(VkDeviceAddress aabbAddress, std::uint32_t count, bool opaque, VkCommandBuffer cmd, RTBuildOptions options = {}, std::uint32_t stride = sizeof(RTAabb));

        // Refit the triangle BLAS against new geometry of the *same
        // topology* (same vertex count and index count as the original
        // Build). When that Build set RTBuildOptions::allowUpdate this issues
        // a hardware UPDATE-mode build — much cheaper than a rebuild, and the
        // BLAS handle / blasAddr are preserved, so TLAS instances referencing
        // it stay valid. A hardware update may only move vertex positions,
        // not change connectivity, so on the UPDATE path only the vertex
        // buffer is re-uploaded; the index buffer is left untouched (its
        // contents are immutable for an UPDATE). The `indicies` span is read
        // only for its count — a bitwise-different but topologically-equivalent
        // index array is ignored on refit. Without allowUpdate (or if the
        // counts changed) it falls back to a full rebuild, which does re-upload
        // both buffers. Lifetime contract matches Build: the spans need only
        // outlive this call. Call this per frame to track a deforming mesh.
        void Refit(std::span<Vector<float, 3, 3>> verticies, std::span<std::uint32_t> indicies, VkCommandBuffer cmd);
        // Procedural analog of Refit: new object-space boxes, same count.
        void RefitProcedural(std::span<const RTAabb> aabbs, VkCommandBuffer cmd);
        // Zero-copy procedural refit: the device-buffer counterpart of
        // RefitProcedural above and the input-source pairing for the device
        // BuildProcedural. The GPU producer rewrote the boxes in `aabbAddress`
        // in place; this re-records the build straight from that address —
        // an in-place hardware UPDATE when the original build set allowUpdate
        // and `count` is unchanged (handle / blasAddr preserved), otherwise a
        // full rebuild. Same buffer-usage / barrier / lifetime contract as the
        // device BuildProcedural. Call per frame to track GPU-written boxes
        // with zero host involvement.
        void RefitProcedural(VkDeviceAddress aabbAddress, std::uint32_t count, VkCommandBuffer cmd, std::uint32_t stride = sizeof(RTAabb));
    };
}
#endif // !CRAFTER_GRAPHICS_WINDOW_DOM

#ifdef CRAFTER_GRAPHICS_WINDOW_DOM
import std;
import Crafter.Math;
import Crafter.Asset;
import :WebGPU;

export namespace Crafter {
    // Software-RT BLAS node, packed to 32 bytes. Matches the WGSL
    // `BVHNode` struct in the RT WGSL prelude (additional/dom-webgpu.js,
    // rtWgslPrelude) byte-for-byte.
    //
    //   primCount == 0          → inner node, children at indices
    //                             firstChildOrPrim and firstChildOrPrim+1.
    //   primCount  > 0          → leaf, `primCount` primitives starting at
    //                             primIndex `firstChildOrPrim` in the
    //                             global primRemap heap.
    //
    // SAH-built BVH2; constructed CPU-side at Build() time, never refit.
    struct BVHNode {
        float         aabbMin[3];
        std::uint32_t firstChildOrPrim;
        float         aabbMax[3];
        std::uint32_t primCount;
    };
    static_assert(sizeof(BVHNode) == 32);

    // One procedural primitive's axis-aligned bounding box, in object
    // space. The analog of VkAabbPositionsKHR — the BLAS stores these
    // instead of triangles and an intersection shader (registered in the
    // hit group as a ProceduralHitGroup) reports the actual surface hit
    // for each AABB the ray enters.
    struct RTAabb {
        float min[3];
        float max[3];
    };
    static_assert(sizeof(RTAabb) == 24);

    // Mirror of the native build-option types so portable code compiles
    // unchanged. The software-RT path has no hardware acceleration
    // structure, so the FastTrace/FastBuild preference has no effect (the
    // SAH BVH2 is always built the same way) and a "refit" simply rebuilds
    // the BVH on the host. The types exist purely for API symmetry.
    enum class RTBuildPreference {
        FastTrace,
        FastBuild,
    };
    struct RTBuildOptions {
        RTBuildPreference preference = RTBuildPreference::FastTrace;
        bool allowUpdate = false;
    };

    class Mesh {
        public:
        // BLAS "handle": opaque identity that goes into
        // RTInstance::accelerationStructureReference. Set by Build() to a
        // stable u32 (widened to u64 for Vulkan-struct layout parity), used
        // by the WebGPU TLAS-build compute shader to look up the BLAS root
        // AABB and per-mesh heap offsets. Handle 0 is the unassigned
        // sentinel; never returned by Build().
        std::uint64_t blasAddr = 0;
        std::uint32_t triangleCount = 0;
        std::uint32_t vertexCount   = 0;

        bool opaque = true;

        // Build BLAS from raw triangle data. Runs the CPU SAH BVH2 builder
        // and forwards vertex/index/BVH/remap arrays to the JS-side mesh
        // heap (additional/dom-webgpu.js), which queue.writeBuffers them
        // into the global heaps and records the per-mesh offsets keyed by
        // the returned handle. The `cmd` parameter is unused on WebGPU —
        // kept for API symmetry with the Vulkan signature.
        void Build(std::span<Crafter::Vector<float, 3, 3>> vertices,
                   std::span<std::uint32_t>                indices,
                   WebGPUCommandEncoderRef                  cmd     = 0,
                   RTBuildOptions                           options = {});

        // CPU-decompress the .cmesh blob (no VK_EXT_memory_decompression
        // equivalent in WebGPU) and forward to the positions+indices path,
        // plus push the optional `data` region into the per-vertex attribs
        // heap so closest-hit shaders can sample UVs / normals / tangents.
        // The data layout is example-defined — the heap is exposed in WGSL
        // as `vertexAttribs : array<u32>` with a per-mesh u32-word offset.
        void Build(const ::Crafter::CompressedMeshAsset& asset,
                   WebGPUCommandEncoderRef                cmd     = 0,
                   RTBuildOptions                         options = {});

        // Build an AABB (procedural) BLAS from a list of object-space boxes
        // — the WebGPU analog of a VK_GEOMETRY_TYPE_AABBS_KHR geometry. The
        // hit group bound to instances of this mesh must be a
        // ProceduralHitGroup carrying an intersection shader; that shader is
        // invoked for each box the ray enters and reports the surface hit.
        // `opaque` is the geometry's opaque bit: pass false to let any-hit
        // shaders run (the default for procedural geometry, which is usually
        // transparent / volumetric). The `cmd` parameter is unused on
        // WebGPU — kept for API symmetry with the triangle path.
        void BuildProcedural(std::span<const RTAabb> aabbs,
                             bool                    opaque  = false,
                             WebGPUCommandEncoderRef cmd     = 0,
                             RTBuildOptions          options = {});

        // Zero-copy procedural build: the boxes already live in an existing
        // device buffer (`aabbBuffer`, a GPUBuffer a compute pass wrote), fed
        // straight into the BLAS with no host copy — the WebGPU analog of the
        // device-address Vulkan BuildProcedural. The software path has no
        // hardware AS and builds its BVH on the host, but here there is no
        // host copy to build over, so the device boxes are copied GPU→GPU into
        // the mesh heap and wrapped in a single root leaf bounded by
        // `worldBounds` (object-space min/max enclosing all `count` boxes —
        // the producer supplies it since the CPU never sees the boxes). The
        // traversal then linearly intersects every box. `opaque` is the
        // geometry opaque bit (false lets any-hit run). `options` has no effect
        // (no hardware AS to tune), kept for API symmetry.
        void BuildProcedural(WebGPUBufferRef         aabbBuffer,
                             std::uint32_t           count,
                             RTAabb                  worldBounds,
                             bool                    opaque  = false,
                             WebGPUCommandEncoderRef cmd     = 0,
                             RTBuildOptions          options = {});

        // Refit analogs of the native API. With no hardware AS to update,
        // these simply re-run the host BVH build over the new data, so they
        // accept the full geometry (not just moved positions). Provided so
        // portable code that calls Refit/RefitProcedural compiles and
        // behaves correctly on the WebGPU backend.
        void Refit(std::span<Crafter::Vector<float, 3, 3>> vertices,
                   std::span<std::uint32_t>                indices,
                   WebGPUCommandEncoderRef                 cmd = 0);
        void RefitProcedural(std::span<const RTAabb> aabbs,
                             WebGPUCommandEncoderRef cmd = 0);
        // Zero-copy procedural refit: re-copy the boxes from the device buffer
        // into this mesh's heap region (count unchanged) and update the root
        // leaf bounds, keeping blasAddr stable — unlike the host-span
        // RefitProcedural, which rebuilds and re-publishes a NEW handle. The
        // device counterpart of the device-address Vulkan RefitProcedural;
        // call per frame to track GPU-written boxes with no host copy.
        void RefitProcedural(WebGPUBufferRef         aabbBuffer,
                             std::uint32_t           count,
                             RTAabb                  worldBounds,
                             WebGPUCommandEncoderRef cmd = 0);
    };
}
#endif // CRAFTER_GRAPHICS_WINDOW_DOM