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

// DOM-mode Mesh implementation: SAH BVH2 built on the host, then
// forwarded to the JS bridge which appends the four data streams
// (vertices, indices, BVH nodes, primRemap) into the global RT mesh
// heaps. The handle returned by wgpuRegisterMeshBLAS goes into
// RTInstance::accelerationStructureReference and lets the TLAS-build
// compute pass and the traversal kernel find the BLAS data later.
//
// BVH layout must stay binary-identical to the WGSL `BVHNode` struct
// declared in additional/dom-webgpu.js (rtWgslPrelude).

module;
module Crafter.Graphics:Mesh_implWebGPU;

import :Mesh;
import :WebGPU;
import Crafter.Asset;
import Crafter.Math;
import std;

using namespace Crafter;

namespace {
    // ─── BVH builder (binned SAH, 8 bins, BVH2) ────────────────────────

    constexpr std::uint32_t kBinCount      = 8;
    constexpr std::uint32_t kMaxLeafSize   = 4;
    constexpr float         kTraversalCost = 1.0f;
    constexpr float         kIntersectCost = 1.0f;

    struct AABB {
        float lo[3] { std::numeric_limits<float>::infinity(),
                      std::numeric_limits<float>::infinity(),
                      std::numeric_limits<float>::infinity() };
        float hi[3] {-std::numeric_limits<float>::infinity(),
                     -std::numeric_limits<float>::infinity(),
                     -std::numeric_limits<float>::infinity() };

        void Extend(const float p[3]) noexcept {
            for (int a = 0; a < 3; ++a) {
                if (p[a] < lo[a]) lo[a] = p[a];
                if (p[a] > hi[a]) hi[a] = p[a];
            }
        }
        void Extend(const AABB& o) noexcept {
            for (int a = 0; a < 3; ++a) {
                if (o.lo[a] < lo[a]) lo[a] = o.lo[a];
                if (o.hi[a] > hi[a]) hi[a] = o.hi[a];
            }
        }
        float SurfaceArea() const noexcept {
            float dx = hi[0] - lo[0];
            float dy = hi[1] - lo[1];
            float dz = hi[2] - lo[2];
            if (dx < 0.0f || dy < 0.0f || dz < 0.0f) return 0.0f;
            return 2.0f * (dx*dy + dx*dz + dy*dz);
        }
    };

    struct PrimRef {
        AABB box;
        float centroid[3];
        std::uint32_t triIndex;
    };

    struct Bin {
        AABB box;
        std::uint32_t count = 0;
    };

    struct Builder {
        std::vector<PrimRef> prims;
        std::vector<BVHNode> nodes;

        std::pair<std::uint32_t, std::uint32_t> AllocateChildren() {
            std::uint32_t l = static_cast<std::uint32_t>(nodes.size());
            nodes.emplace_back();
            nodes.emplace_back();
            return { l, l + 1 };
        }

        void BuildRecursive(std::uint32_t nodeIdx,
                            std::uint32_t first,
                            std::uint32_t count) {
            AABB bounds, centroidBounds;
            for (std::uint32_t i = 0; i < count; ++i) {
                const auto& p = prims[first + i];
                bounds.Extend(p.box);
                centroidBounds.Extend(p.centroid);
            }

            auto emitLeaf = [&] {
                BVHNode& n = nodes[nodeIdx];
                std::memcpy(n.aabbMin, bounds.lo, sizeof(bounds.lo));
                std::memcpy(n.aabbMax, bounds.hi, sizeof(bounds.hi));
                n.firstChildOrPrim = first;
                n.primCount        = count;
            };

            if (count <= kMaxLeafSize) { emitLeaf(); return; }

            int   bestAxis  = -1;
            float bestCost  = std::numeric_limits<float>::infinity();
            std::uint32_t bestBin = 0;

            float parentArea = bounds.SurfaceArea();
            if (parentArea <= 0.0f) { emitLeaf(); return; }

            for (int axis = 0; axis < 3; ++axis) {
                float extent = centroidBounds.hi[axis] - centroidBounds.lo[axis];
                if (extent <= 0.0f) continue;
                float invExtent = static_cast<float>(kBinCount) / extent;

                std::array<Bin, kBinCount> bins{};
                for (std::uint32_t i = 0; i < count; ++i) {
                    const auto& p = prims[first + i];
                    float t = (p.centroid[axis] - centroidBounds.lo[axis]) * invExtent;
                    std::uint32_t b = static_cast<std::uint32_t>(t);
                    if (b >= kBinCount) b = kBinCount - 1;
                    bins[b].box.Extend(p.box);
                    bins[b].count += 1;
                }

                std::array<AABB,         kBinCount - 1> leftBox;
                std::array<std::uint32_t,kBinCount - 1> leftCount{};
                {
                    AABB acc; std::uint32_t cnt = 0;
                    for (std::uint32_t i = 0; i < kBinCount - 1; ++i) {
                        acc.Extend(bins[i].box);
                        cnt += bins[i].count;
                        leftBox[i]   = acc;
                        leftCount[i] = cnt;
                    }
                }
                {
                    AABB acc; std::uint32_t cnt = 0;
                    for (std::int32_t i = kBinCount - 1; i >= 1; --i) {
                        acc.Extend(bins[i].box);
                        cnt += bins[i].count;
                        std::uint32_t split = static_cast<std::uint32_t>(i - 1);
                        if (leftCount[split] == 0 || cnt == 0) continue;
                        float cost = kTraversalCost
                                   + (leftBox[split].SurfaceArea() * leftCount[split]
                                      + acc.SurfaceArea() * cnt) * kIntersectCost / parentArea;
                        if (cost < bestCost) {
                            bestCost = cost;
                            bestAxis = axis;
                            bestBin  = split;
                        }
                    }
                }
            }

            float leafCost = static_cast<float>(count) * kIntersectCost;
            if (bestAxis < 0 || bestCost >= leafCost) { emitLeaf(); return; }

            float invExtent = static_cast<float>(kBinCount)
                            / (centroidBounds.hi[bestAxis] - centroidBounds.lo[bestAxis]);
            float lo = centroidBounds.lo[bestAxis];
            auto mid = std::partition(
                prims.begin() + first, prims.begin() + first + count,
                [&](const PrimRef& p) {
                    float t = (p.centroid[bestAxis] - lo) * invExtent;
                    std::uint32_t b = static_cast<std::uint32_t>(t);
                    if (b >= kBinCount) b = kBinCount - 1;
                    return b <= bestBin;
                });
            std::uint32_t leftCount =
                static_cast<std::uint32_t>(mid - (prims.begin() + first));
            if (leftCount == 0 || leftCount == count) { emitLeaf(); return; }

            auto [leftIdx, rightIdx] = AllocateChildren();
            {
                BVHNode& n = nodes[nodeIdx];
                std::memcpy(n.aabbMin, bounds.lo, sizeof(bounds.lo));
                std::memcpy(n.aabbMax, bounds.hi, sizeof(bounds.hi));
                n.firstChildOrPrim = leftIdx;
                n.primCount        = 0;
            }
            BuildRecursive(leftIdx,  first,             leftCount);
            BuildRecursive(rightIdx, first + leftCount, count - leftCount);
        }

        void Build(std::span<const Vector<float, 3, 3>> vertices,
                   std::span<const std::uint32_t>       indices) {
            std::uint32_t triCount = static_cast<std::uint32_t>(indices.size()) / 3;
            prims.resize(triCount);
            for (std::uint32_t i = 0; i < triCount; ++i) {
                std::uint32_t i0 = indices[i*3 + 0];
                std::uint32_t i1 = indices[i*3 + 1];
                std::uint32_t i2 = indices[i*3 + 2];
                const auto& v0 = vertices[i0];
                const auto& v1 = vertices[i1];
                const auto& v2 = vertices[i2];
                float p0[3] { v0.v[0], v0.v[1], v0.v[2] };
                float p1[3] { v1.v[0], v1.v[1], v1.v[2] };
                float p2[3] { v2.v[0], v2.v[1], v2.v[2] };
                auto& pr = prims[i];
                pr.box.Extend(p0);
                pr.box.Extend(p1);
                pr.box.Extend(p2);
                pr.centroid[0] = (pr.box.lo[0] + pr.box.hi[0]) * 0.5f;
                pr.centroid[1] = (pr.box.lo[1] + pr.box.hi[1]) * 0.5f;
                pr.centroid[2] = (pr.box.lo[2] + pr.box.hi[2]) * 0.5f;
                pr.triIndex = i;
            }
            nodes.reserve(triCount * 2);
            nodes.emplace_back();
            BuildRecursive(0, 0, triCount);
        }

        // AABB (procedural) geometry: one PrimRef per box, the box itself
        // is the primitive bound. The same SAH BVH2 then partitions them.
        void BuildFromAabbs(std::span<const RTAabb> aabbs) {
            std::uint32_t count = static_cast<std::uint32_t>(aabbs.size());
            prims.resize(count);
            for (std::uint32_t i = 0; i < count; ++i) {
                auto& pr = prims[i];
                pr.box.Extend(aabbs[i].min);
                pr.box.Extend(aabbs[i].max);
                pr.centroid[0] = (pr.box.lo[0] + pr.box.hi[0]) * 0.5f;
                pr.centroid[1] = (pr.box.lo[1] + pr.box.hi[1]) * 0.5f;
                pr.centroid[2] = (pr.box.lo[2] + pr.box.hi[2]) * 0.5f;
                pr.triIndex = i;
            }
            nodes.reserve(count * 2);
            nodes.emplace_back();
            BuildRecursive(0, 0, count);
        }
    };
}

namespace {
    // Shared between the positions-only and the compressed-asset Build paths.
    // attribsBytes is empty for positions-only meshes; the JS bridge skips
    // the attribs-heap append in that case.
    void BuildBVHAndRegister(Mesh& mesh,
                             std::span<const Vector<float, 3, 3>> vertices,
                             std::span<const std::uint32_t>       indices,
                             std::span<const std::byte>           attribsBytes) {
        mesh.triangleCount = static_cast<std::uint32_t>(indices.size()) / 3;
        mesh.vertexCount   = static_cast<std::uint32_t>(vertices.size());

        Builder builder;
        builder.Build(vertices, indices);

        std::vector<std::uint32_t> primRemap(mesh.triangleCount);
        for (std::uint32_t i = 0; i < mesh.triangleCount; ++i) {
            primRemap[i] = builder.prims[i].triIndex;
        }

        const BVHNode& root = builder.nodes[0];
        mesh.blasAddr = WebGPU::wgpuRegisterMeshBLAS(
            root.aabbMin[0], root.aabbMin[1], root.aabbMin[2],
            root.aabbMax[0], root.aabbMax[1], root.aabbMax[2],
            vertices.data(),       static_cast<std::int32_t>(vertices.size()),
            indices.data(),        static_cast<std::int32_t>(indices.size()),
            builder.nodes.data(),  static_cast<std::int32_t>(builder.nodes.size()),
            primRemap.data(),      static_cast<std::int32_t>(primRemap.size()),
            attribsBytes.data(),   static_cast<std::int32_t>(attribsBytes.size()),
            /*geomType*/ 0,
            /*opaqueFlag*/ mesh.opaque ? 1 : 0,
            /*primCount*/ static_cast<std::int32_t>(mesh.triangleCount));
    }
}

void Mesh::Build(std::span<Vector<float, 3, 3>> vertices,
                 std::span<std::uint32_t>       indices,
                 WebGPUCommandEncoderRef        /*cmd*/,
                 RTBuildOptions                 /*options*/) {
    // The build preference (FastTrace/FastBuild) and allowUpdate have no
    // effect on the software-RT path — there is no hardware AS to tune or
    // refit. The SAH BVH2 is always built the same way.
    BuildBVHAndRegister(*this, vertices, indices, {});
}

void Mesh::Build(const CompressedMeshAsset& asset,
                 WebGPUCommandEncoderRef    /*cmd*/,
                 RTBuildOptions             /*options*/) {
    std::vector<Vector<float, 3, 3>> vertices(asset.vertexCount);
    std::vector<std::uint32_t>       indices(asset.indexCount);
    std::vector<std::byte>           dataBytes(
        static_cast<std::size_t>(asset.dataCount) * asset.dataStride);

    // CompressedBlob always carries 3 regions for MeshAsset (the data region
    // can have decompressedSize=0). DecompressCPU validates output sizes
    // against region sizes, so the empty-data path needs the empty span.
    std::array<std::span<std::byte>, 3> outputs = {
        std::as_writable_bytes(std::span(vertices)),
        std::as_writable_bytes(std::span(indices)),
        std::span<std::byte>(dataBytes),
    };
    Compression::DecompressCPU(asset.blob,
        std::span(outputs).first(asset.blob.regions.size()));

    BuildBVHAndRegister(*this, vertices, indices, std::span(dataBytes));
}

void Mesh::BuildProcedural(std::span<const RTAabb> aabbs,
                           bool                    opaque_,
                           WebGPUCommandEncoderRef /*cmd*/,
                           RTBuildOptions          /*options*/) {
    const std::uint32_t count = static_cast<std::uint32_t>(aabbs.size());
    opaque        = opaque_;
    triangleCount = 0;                 // not a triangle mesh
    vertexCount   = count * 2;         // 2 "vertices" (min,max) per box

    Builder builder;
    builder.BuildFromAabbs(aabbs);

    // The AABB stream is uploaded in *original* primitive order (2 vec3 per
    // box). primRemap maps each BVH leaf slot back to its original index, so
    // the intersection shader's _rtFetchAabb(meshRec, primId) reads the
    // right box — exactly mirroring how the triangle path indexes vertices.
    std::vector<Vector<float, 3, 3>> boxVerts(count * 2);
    for (std::uint32_t i = 0; i < count; ++i) {
        boxVerts[i*2 + 0] = Vector<float, 3, 3>{ aabbs[i].min[0], aabbs[i].min[1], aabbs[i].min[2] };
        boxVerts[i*2 + 1] = Vector<float, 3, 3>{ aabbs[i].max[0], aabbs[i].max[1], aabbs[i].max[2] };
    }

    std::vector<std::uint32_t> primRemap(count);
    for (std::uint32_t i = 0; i < count; ++i) {
        primRemap[i] = builder.prims[i].triIndex;
    }

    const BVHNode& root = builder.nodes[0];
    blasAddr = WebGPU::wgpuRegisterMeshBLAS(
        root.aabbMin[0], root.aabbMin[1], root.aabbMin[2],
        root.aabbMax[0], root.aabbMax[1], root.aabbMax[2],
        boxVerts.data(),      static_cast<std::int32_t>(boxVerts.size()),
        nullptr,              0,
        builder.nodes.data(), static_cast<std::int32_t>(builder.nodes.size()),
        primRemap.data(),     static_cast<std::int32_t>(primRemap.size()),
        nullptr,              0,
        /*geomType*/ 1,
        /*opaqueFlag*/ opaque ? 1 : 0,
        /*primCount*/ static_cast<std::int32_t>(count));
}

void Mesh::BuildProcedural(WebGPUBufferRef         aabbBuffer,
                           std::uint32_t           count,
                           RTAabb                  worldBounds,
                           bool                    opaque_,
                           WebGPUCommandEncoderRef /*cmd*/,
                           RTBuildOptions          /*options*/) {
    // Zero-copy: the boxes are already on the GPU (a compute pass wrote
    // them). The bridge copies them GPU→GPU into the vertices heap and
    // registers a single-root-leaf BLAS bounded by worldBounds — no SAH
    // build (there is no host copy of the boxes to build over) and no host
    // round-trip. Traversal linearly intersects all `count` boxes.
    opaque        = opaque_;
    triangleCount = 0;             // not a triangle mesh
    vertexCount   = count * 2;     // 2 "vertices" (min,max) per box

    blasAddr = WebGPU::wgpuRegisterMeshBLASDeviceAabbs(
        aabbBuffer, static_cast<std::int32_t>(count),
        worldBounds.min[0], worldBounds.min[1], worldBounds.min[2],
        worldBounds.max[0], worldBounds.max[1], worldBounds.max[2],
        opaque ? 1 : 0);
}

void Mesh::Refit(std::span<Vector<float, 3, 3>> vertices,
                 std::span<std::uint32_t>       indices,
                 WebGPUCommandEncoderRef        cmd) {
    // No hardware AS to update in place — the software path rebuilds the
    // host BVH and registers it afresh. Unlike the Vulkan UPDATE path, this
    // assigns a NEW blasAddr (the JS heap append is not in-place), so any
    // RTInstance::accelerationStructureReference pointing at this mesh must
    // be re-pointed at the updated mesh.blasAddr afterwards. Refit is far
    // better suited to the hardware backend; on WebGPU prefer rebuilding
    // and re-publishing the handle.
    Build(vertices, indices, cmd);
}

void Mesh::RefitProcedural(std::span<const RTAabb> aabbs,
                           WebGPUCommandEncoderRef cmd) {
    BuildProcedural(aabbs, opaque, cmd);
}

void Mesh::RefitProcedural(WebGPUBufferRef         aabbBuffer,
                           std::uint32_t           count,
                           RTAabb                  worldBounds,
                           WebGPUCommandEncoderRef /*cmd*/) {
    // Re-copy the GPU boxes into the existing heap region and refresh the
    // root bounds — keeps blasAddr stable (no re-register), so TLAS
    // instances referencing this mesh stay valid across the per-frame
    // update. No host copy.
    vertexCount = count * 2;
    WebGPU::wgpuRefitMeshBLASDeviceAabbs(
        static_cast<std::uint32_t>(blasAddr), aabbBuffer,
        static_cast<std::int32_t>(count),
        worldBounds.min[0], worldBounds.min[1], worldBounds.min[2],
        worldBounds.max[0], worldBounds.max[1], worldBounds.max[2]);
}
