fn_registry/cpp/functions/gfx/gltf_load_mesh.cpp

#include "gfx/gltf_load_mesh.h"

#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <string>
#include <vector>

// nlohmann/json vendored
#include "nlohmann/json.hpp"

namespace fn::gfx {

// ---------------------------------------------------------------------------
// Thread-local last error
// ---------------------------------------------------------------------------

static thread_local char s_last_error[512] = "";

static void set_error(const char* msg) {
    std::strncpy(s_last_error, msg, sizeof(s_last_error) - 1);
    s_last_error[sizeof(s_last_error) - 1] = '\0';
}

const char* gltf_load_last_error() { return s_last_error; }

// ---------------------------------------------------------------------------
// GLB binary format constants (spec glTF 2.0)
// ---------------------------------------------------------------------------

static constexpr uint32_t GLB_MAGIC        = 0x46546C67u; // "glTF"
static constexpr uint32_t GLB_VERSION      = 2u;
static constexpr uint32_t CHUNK_JSON       = 0x4E4F534Au; // "JSON"
static constexpr uint32_t CHUNK_BIN        = 0x004E4942u; // "BIN\0"

// glTF accessor componentType
static constexpr int CT_UNSIGNED_BYTE  = 5121;
static constexpr int CT_UNSIGNED_SHORT = 5123;
static constexpr int CT_UNSIGNED_INT   = 5125;
static constexpr int CT_FLOAT          = 5126;

// ---------------------------------------------------------------------------
// Math helpers
// ---------------------------------------------------------------------------

static void cross3(const float a[3], const float b[3], float out[3]) {
    out[0] = a[1]*b[2] - a[2]*b[1];
    out[1] = a[2]*b[0] - a[0]*b[2];
    out[2] = a[0]*b[1] - a[1]*b[0];
}

static float dot3(const float a[3], const float b[3]) {
    return a[0]*b[0] + a[1]*b[1] + a[2]*b[2];
}

static float len3(const float a[3]) {
    return std::sqrt(dot3(a, a));
}

// Multiply 4x4 column-major matrix by vec3 (point, w=1)
static void mat4_mul_point(const float m[16], const float p[3], float out[3]) {
    out[0] = m[0]*p[0] + m[4]*p[1] + m[8] *p[2] + m[12];
    out[1] = m[1]*p[0] + m[5]*p[1] + m[9] *p[2] + m[13];
    out[2] = m[2]*p[0] + m[6]*p[1] + m[10]*p[2] + m[14];
}

// Multiply 4x4 column-major matrix by vec3 (direction, w=0 — for normals use
// inverse-transpose, which here is computed at call site)
static void mat4_mul_dir(const float m[16], const float v[3], float out[3]) {
    out[0] = m[0]*v[0] + m[4]*v[1] + m[8] *v[2];
    out[1] = m[1]*v[0] + m[5]*v[1] + m[9] *v[2];
    out[2] = m[2]*v[0] + m[6]*v[1] + m[10]*v[2];
}

// 3x3 inverse-transpose (for normal transform) extracted from upper-left of 4x4.
// Returns false if matrix is singular (scale 0).
static bool compute_normal_matrix(const float m[16], float out[9]) {
    // Extract upper-left 3x3 (column-major from 4x4)
    float a00=m[0], a10=m[1], a20=m[2];
    float a01=m[4], a11=m[5], a21=m[6];
    float a02=m[8], a12=m[9], a22=m[10];

    float det = a00*(a11*a22 - a21*a12)
              - a01*(a10*a22 - a20*a12)
              + a02*(a10*a21 - a20*a11);
    if (std::fabs(det) < 1e-12f) return false;
    float inv = 1.0f / det;

    // Inverse of 3x3, then transpose → inverse-transpose columns become rows
    out[0] = inv * (a11*a22 - a21*a12);
    out[1] = inv * (a21*a02 - a01*a22);
    out[2] = inv * (a01*a12 - a11*a02);
    out[3] = inv * (a20*a12 - a10*a22);
    out[4] = inv * (a00*a22 - a20*a02);
    out[5] = inv * (a10*a02 - a00*a12);
    out[6] = inv * (a10*a21 - a20*a11);
    out[7] = inv * (a20*a01 - a00*a21);
    out[8] = inv * (a00*a11 - a10*a01);
    return true;
}

static void nrm3x3_mul(const float m[9], const float v[3], float out[3]) {
    out[0] = m[0]*v[0] + m[3]*v[1] + m[6]*v[2];
    out[1] = m[1]*v[0] + m[4]*v[1] + m[7]*v[2];
    out[2] = m[2]*v[0] + m[5]*v[1] + m[8]*v[2];
}

// TRS → column-major 4x4 matrix
// translation=[tx,ty,tz], rotation quaternion=[qx,qy,qz,qw], scale=[sx,sy,sz]
static void trs_to_mat4(const float t[3], const float q[4], const float s[3],
                         float out[16]) {
    float qx=q[0], qy=q[1], qz=q[2], qw=q[3];
    float x2=qx+qx, y2=qy+qy, z2=qz+qz;
    float xx=qx*x2, xy=qx*y2, xz=qx*z2;
    float yy=qy*y2, yz=qy*z2, zz=qz*z2;
    float wx=qw*x2, wy=qw*y2, wz=qw*z2;

    out[0]  = (1-(yy+zz))*s[0]; out[1]  = (xy+wz)*s[0];    out[2]  = (xz-wy)*s[0];    out[3]  = 0;
    out[4]  = (xy-wz)*s[1];    out[5]  = (1-(xx+zz))*s[1]; out[6]  = (yz+wx)*s[1];    out[7]  = 0;
    out[8]  = (xz+wy)*s[2];    out[9]  = (yz-wx)*s[2];    out[10] = (1-(xx+yy))*s[2]; out[11] = 0;
    out[12] = t[0];             out[13] = t[1];             out[14] = t[2];              out[15] = 1;
}

// ---------------------------------------------------------------------------
// Accessor reading helpers
// ---------------------------------------------------------------------------

struct BufView {
    const uint8_t* base  = nullptr;
    size_t         total = 0;
};

// Read a single element of 'count' components from accessor at element index 'idx'.
// component_type: CT_FLOAT, CT_UNSIGNED_BYTE, CT_UNSIGNED_SHORT, CT_UNSIGNED_INT
// components_per_element: 1 (SCALAR) or 3 (VEC3) etc.
// Returns false if out-of-bounds.
static bool read_float_vec(const BufView& bin,
                            int component_type,
                            int components_per_element,
                            size_t byte_offset,  // accessor.byteOffset + bufferView.byteOffset
                            int byte_stride,     // bufferView.byteStride (0 = tightly packed)
                            size_t idx,
                            float out[4]) {
    size_t comp_size = 0;
    switch (component_type) {
        case CT_UNSIGNED_BYTE:  comp_size = 1; break;
        case CT_UNSIGNED_SHORT: comp_size = 2; break;
        case CT_UNSIGNED_INT:   comp_size = 4; break;
        case CT_FLOAT:          comp_size = 4; break;
        default: return false;
    }
    size_t element_size = comp_size * (size_t)components_per_element;
    size_t stride = (byte_stride > 0) ? (size_t)byte_stride : element_size;
    size_t off    = byte_offset + idx * stride;
    if (off + element_size > bin.total) return false;

    const uint8_t* p = bin.base + off;
    for (int c = 0; c < components_per_element; ++c) {
        const uint8_t* cp = p + (size_t)c * comp_size;
        switch (component_type) {
            case CT_UNSIGNED_BYTE:  out[c] = (float)*cp; break;
            case CT_UNSIGNED_SHORT: {
                uint16_t v; std::memcpy(&v, cp, 2); out[c] = (float)v; break;
            }
            case CT_UNSIGNED_INT: {
                uint32_t v; std::memcpy(&v, cp, 4); out[c] = (float)v; break;
            }
            case CT_FLOAT: {
                float v; std::memcpy(&v, cp, 4); out[c] = v; break;
            }
            default: return false;
        }
    }
    return true;
}

static bool read_index(const BufView& bin,
                        int component_type,
                        size_t byte_offset,
                        size_t idx,
                        uint32_t& out) {
    float v[1] = {};
    if (!read_float_vec(bin, component_type, 1, byte_offset, 0, idx, v))
        return false;
    out = static_cast<uint32_t>(v[0]);
    return true;
}

// ---------------------------------------------------------------------------
// Core GLB parser
// ---------------------------------------------------------------------------

static Mesh parse_glb(const uint8_t* data, size_t size) {
    s_last_error[0] = '\0';

    // --- 1. Validate header (12 bytes) ---
    if (size < 12) { set_error("file too small for GLB header"); return {}; }

    uint32_t magic, version, total_len;
    std::memcpy(&magic,     data,     4);
    std::memcpy(&version,   data + 4, 4);
    std::memcpy(&total_len, data + 8, 4);

    if (magic != GLB_MAGIC)   { set_error("not a GLB file (bad magic)"); return {}; }
    if (version != GLB_VERSION){ set_error("unsupported GLB version (expected 2)"); return {}; }
    if (total_len > size)       { set_error("GLB total_length > buffer size"); return {}; }

    // --- 2. Walk chunks ---
    const uint8_t* json_data  = nullptr; size_t json_len  = 0;
    const uint8_t* bin_data   = nullptr; size_t bin_len   = 0;

    size_t pos = 12;
    while (pos + 8 <= total_len) {
        uint32_t chunk_len, chunk_type;
        std::memcpy(&chunk_len,  data + pos,     4);
        std::memcpy(&chunk_type, data + pos + 4, 4);
        pos += 8;
        if (pos + chunk_len > total_len) { set_error("chunk extends past file end"); return {}; }
        if (chunk_type == CHUNK_JSON) {
            json_data = data + pos;
            json_len  = chunk_len;
        } else if (chunk_type == CHUNK_BIN) {
            bin_data = data + pos;
            bin_len  = chunk_len;
        }
        pos += chunk_len;
    }

    if (!json_data) { set_error("no JSON chunk found"); return {}; }

    // --- 3. Parse JSON ---
    nlohmann::json j;
    try {
        j = nlohmann::json::parse(json_data, json_data + json_len);
    } catch (const std::exception& e) {
        std::snprintf(s_last_error, sizeof(s_last_error), "JSON parse error: %s", e.what());
        return {};
    }

    // --- 4. Find first mesh / first primitive ---
    if (!j.contains("meshes") || j["meshes"].empty()) {
        set_error("no meshes in glTF");
        return {};
    }
    auto& prim = j["meshes"][0]["primitives"][0];

    auto& attrs = prim["attributes"];
    if (!attrs.contains("POSITION")) {
        set_error("primitive has no POSITION attribute");
        return {};
    }

    auto& accessors   = j["accessors"];
    auto& bufferViews = j["bufferViews"];

    BufView bin_view { bin_data, bin_len };

    // Helper: resolve accessor index → (byte_offset, byte_stride, component_type, count, components_per_elem)
    struct AccInfo { size_t byte_offset; int byte_stride; int comp_type; size_t count; int ncomp; };

    auto resolve_accessor = [&](int acc_idx, AccInfo& out) -> bool {
        if (acc_idx < 0 || acc_idx >= (int)accessors.size()) return false;
        auto& acc = accessors[acc_idx];
        int bv_idx = acc.value("bufferView", -1);
        size_t acc_offset = acc.value("byteOffset", 0);
        out.comp_type = acc.value("componentType", 0);
        out.count     = acc.value("count", 0u);
        std::string type_str = acc.value("type", "SCALAR");
        out.ncomp = 1;
        if (type_str == "VEC2") out.ncomp = 2;
        else if (type_str == "VEC3") out.ncomp = 3;
        else if (type_str == "VEC4") out.ncomp = 4;

        if (bv_idx >= 0 && bv_idx < (int)bufferViews.size()) {
            auto& bv = bufferViews[bv_idx];
            size_t bv_offset = bv.value("byteOffset", 0u);
            out.byte_stride  = bv.value("byteStride", 0);
            out.byte_offset  = acc_offset + bv_offset;
        } else {
            out.byte_offset = acc_offset;
            out.byte_stride = 0;
        }
        return out.count > 0 && out.comp_type != 0;
    };

    // --- 5. Read POSITION ---
    AccInfo pos_info{};
    if (!resolve_accessor(attrs["POSITION"].get<int>(), pos_info)) {
        set_error("failed to resolve POSITION accessor");
        return {};
    }
    if (pos_info.ncomp != 3 || pos_info.comp_type != CT_FLOAT) {
        set_error("POSITION must be float vec3");
        return {};
    }
    if (!bin_data && pos_info.count > 0) {
        set_error("POSITION accessor requires BIN chunk, which is missing");
        return {};
    }

    size_t nv = pos_info.count;
    std::vector<float> positions(nv * 3);
    for (size_t i = 0; i < nv; ++i) {
        float v[4]{};
        if (!read_float_vec(bin_view, CT_FLOAT, 3, pos_info.byte_offset,
                             pos_info.byte_stride, i, v)) {
            set_error("out-of-bounds read in POSITION");
            return {};
        }
        positions[i*3+0] = v[0];
        positions[i*3+1] = v[1];
        positions[i*3+2] = v[2];
    }

    // --- 6. Read NORMAL (optional) ---
    std::vector<float> normals;
    bool has_normals = false;
    if (attrs.contains("NORMAL")) {
        AccInfo nrm_info{};
        if (resolve_accessor(attrs["NORMAL"].get<int>(), nrm_info) &&
            nrm_info.ncomp == 3 && nrm_info.comp_type == CT_FLOAT &&
            nrm_info.count == nv) {
            normals.resize(nv * 3);
            for (size_t i = 0; i < nv; ++i) {
                float v[4]{};
                if (!read_float_vec(bin_view, CT_FLOAT, 3, nrm_info.byte_offset,
                                     nrm_info.byte_stride, i, v)) {
                    set_error("out-of-bounds read in NORMAL");
                    return {};
                }
                normals[i*3+0] = v[0];
                normals[i*3+1] = v[1];
                normals[i*3+2] = v[2];
            }
            has_normals = true;
        }
    }

    // --- 7. Read indices ---
    std::vector<uint32_t> indices;
    if (prim.contains("indices") && !prim["indices"].is_null()) {
        AccInfo idx_info{};
        int idx_acc = prim["indices"].get<int>();
        if (!resolve_accessor(idx_acc, idx_info)) {
            set_error("failed to resolve indices accessor");
            return {};
        }
        if (!bin_data && idx_info.count > 0) {
            set_error("indices accessor requires BIN chunk, which is missing");
            return {};
        }
        indices.resize(idx_info.count);
        for (size_t i = 0; i < idx_info.count; ++i) {
            if (!read_index(bin_view, idx_info.comp_type, idx_info.byte_offset, i, indices[i])) {
                set_error("out-of-bounds read in indices");
                return {};
            }
        }
    } else {
        // No indices: interpret as sequential triangle list
        indices.resize(nv);
        for (size_t i = 0; i < nv; ++i) indices[i] = (uint32_t)i;
    }

    // --- 8. Generate normals if missing (smooth, area-weighted) ---
    if (!has_normals) {
        normals.assign(nv * 3, 0.0f);
        size_t ntri = indices.size() / 3;
        for (size_t t = 0; t < ntri; ++t) {
            uint32_t i0 = indices[t*3+0];
            uint32_t i1 = indices[t*3+1];
            uint32_t i2 = indices[t*3+2];
            if (i0 >= nv || i1 >= nv || i2 >= nv) continue;

            float e1[3] = {
                positions[i1*3+0] - positions[i0*3+0],
                positions[i1*3+1] - positions[i0*3+1],
                positions[i1*3+2] - positions[i0*3+2]
            };
            float e2[3] = {
                positions[i2*3+0] - positions[i0*3+0],
                positions[i2*3+1] - positions[i0*3+1],
                positions[i2*3+2] - positions[i0*3+2]
            };
            float face_n[3];
            cross3(e1, e2, face_n);
            // face_n magnitude = 2 * area → area weighting automatic
            for (uint32_t vi : {i0, i1, i2}) {
                normals[vi*3+0] += face_n[0];
                normals[vi*3+1] += face_n[1];
                normals[vi*3+2] += face_n[2];
            }
        }
        // Normalize per-vertex
        for (size_t i = 0; i < nv; ++i) {
            float* n = &normals[i*3];
            float l = len3(n);
            if (l > 1e-8f) { n[0]/=l; n[1]/=l; n[2]/=l; }
            else            { n[0]=0; n[1]=1; n[2]=0; } // degenerate fallback
        }
    }

    // --- 9. Apply node transform (first node referencing this mesh) ---
    bool applied_transform = false;
    if (j.contains("nodes") && !j["nodes"].empty()) {
        auto& nodes = j["nodes"];
        for (size_t ni = 0; ni < nodes.size() && !applied_transform; ++ni) {
            auto& node = nodes[ni];
            if (!node.contains("mesh") || node["mesh"].get<int>() != 0) continue;

            float mat[16] = {
                1,0,0,0, 0,1,0,0, 0,0,1,0, 0,0,0,1
            }; // identity column-major

            if (node.contains("matrix") && node["matrix"].is_array() && node["matrix"].size() == 16) {
                for (int k = 0; k < 16; ++k)
                    mat[k] = node["matrix"][k].get<float>();
                applied_transform = true;
            } else {
                float t[3] = {0,0,0}, q[4] = {0,0,0,1}, s[3] = {1,1,1};
                bool has_trs = false;
                if (node.contains("translation") && node["translation"].size() == 3) {
                    for (int k = 0; k < 3; ++k) t[k] = node["translation"][k].get<float>();
                    has_trs = true;
                }
                if (node.contains("rotation") && node["rotation"].size() == 4) {
                    for (int k = 0; k < 4; ++k) q[k] = node["rotation"][k].get<float>();
                    has_trs = true;
                }
                if (node.contains("scale") && node["scale"].size() == 3) {
                    for (int k = 0; k < 3; ++k) s[k] = node["scale"][k].get<float>();
                    has_trs = true;
                }
                if (has_trs) {
                    trs_to_mat4(t, q, s, mat);
                    applied_transform = true;
                }
            }

            if (applied_transform) {
                // Check if matrix is non-trivially identity
                const float id[16] = {1,0,0,0, 0,1,0,0, 0,0,1,0, 0,0,0,1};
                bool is_identity = true;
                for (int k = 0; k < 16; ++k)
                    if (std::fabs(mat[k] - id[k]) > 1e-6f) { is_identity = false; break; }

                if (!is_identity) {
                    float nrm_mat[9];
                    bool has_nrm_mat = compute_normal_matrix(mat, nrm_mat);

                    for (size_t vi = 0; vi < nv; ++vi) {
                        float p[3] = { positions[vi*3+0], positions[vi*3+1], positions[vi*3+2] };
                        float tp[3];
                        mat4_mul_point(mat, p, tp);
                        positions[vi*3+0] = tp[0];
                        positions[vi*3+1] = tp[1];
                        positions[vi*3+2] = tp[2];

                        if (has_nrm_mat) {
                            float n[3] = { normals[vi*3+0], normals[vi*3+1], normals[vi*3+2] };
                            float tn[3];
                            nrm3x3_mul(nrm_mat, n, tn);
                            float l = len3(tn);
                            if (l > 1e-8f) { tn[0]/=l; tn[1]/=l; tn[2]/=l; }
                            normals[vi*3+0] = tn[0];
                            normals[vi*3+1] = tn[1];
                            normals[vi*3+2] = tn[2];
                        }
                    }
                }
            }
        }
    }

    Mesh m;
    m.positions = std::move(positions);
    m.normals   = std::move(normals);
    m.indices   = std::move(indices);
    return m;
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

Mesh gltf_load_mesh_from_memory(const unsigned char* data, size_t size) {
    return parse_glb(reinterpret_cast<const uint8_t*>(data), size);
}

Mesh gltf_load_mesh_from_file(const char* path) {
    std::ifstream f(path, std::ios::binary | std::ios::ate);
    if (!f) {
        std::snprintf(s_last_error, sizeof(s_last_error),
                      "cannot open file: %s", path);
        return {};
    }
    auto file_size = f.tellg();
    if (file_size <= 0) { set_error("file is empty"); return {}; }
    f.seekg(0);
    std::vector<uint8_t> buf((size_t)file_size);
    if (!f.read(reinterpret_cast<char*>(buf.data()), file_size)) {
        set_error("file read failed");
        return {};
    }
    return parse_glb(buf.data(), buf.size());
}

} // namespace fn::gfx