fn_registry/cpp/functions/viz/mesh_viewer.cpp

#include "viz/mesh_viewer.h"

#include "gfx/gl_loader.h"
#include "gfx/gl_framebuffer.h"
#include "gfx/gl_shader.h"

#include <string>
#include <unordered_map>

namespace fn::viz {

namespace {

// Vertex+fragment shader with Lambert headlight (light = camera direction).
// We bypass gl_shader::compile_fragment because that helper assumes a
// fullscreen-quad pipeline; here we need explicit attribs and a vertex shader
// that consumes mvp + normalMatrix.

const char* kVert = R"glsl(
#version 330 core
layout(location = 0) in vec3 a_pos;
layout(location = 1) in vec3 a_normal;

uniform mat4 u_view;
uniform mat4 u_proj;

out vec3 v_normal_view;

void main() {
    vec4 view_pos = u_view * vec4(a_pos, 1.0);
    // For pure rotation in the view part of orbit-camera, the upper-left 3x3
    // suffices as the normal matrix.
    v_normal_view = mat3(u_view) * a_normal;
    gl_Position = u_proj * view_pos;
}
)glsl";

const char* kFrag = R"glsl(
#version 330 core
in vec3 v_normal_view;
out vec4 frag;

uniform vec4 u_color;

void main() {
    // Headlight: light dir = +Z in view space (toward camera in right-handed view).
    vec3 N = normalize(v_normal_view);
    vec3 L = vec3(0.0, 0.0, 1.0);
    float ndl = max(dot(N, L), 0.0);
    // ambient 0.2 + diffuse 0.8
    float lighting = 0.2 + 0.8 * ndl;
    frag = vec4(u_color.rgb * lighting, u_color.a);
}
)glsl";

// Per-id cached state.
struct Cache {
    fn::gfx::Framebuffer fb{};
    GLuint               program = 0;
    GLint                loc_view = -1;
    GLint                loc_proj = -1;
    GLint                loc_color = -1;
    bool                 initialized = false;
};

std::unordered_map<std::string, Cache>& caches() {
    static std::unordered_map<std::string, Cache> m;
    return m;
}

GLuint compile_program() {
    // We need vertex+fragment, but gl_shader::compile_fragment only does
    // fragment with a fixed vertex. Inline a small compile here.
    auto compile = [](GLenum type, const char* src) -> GLuint {
        GLuint sh = glCreateShader(type);
        glShaderSource(sh, 1, &src, nullptr);
        glCompileShader(sh);
        GLint ok = 0;
        glGetShaderiv(sh, GL_COMPILE_STATUS, &ok);
        if (!ok) {
            char log[512];
            glGetShaderInfoLog(sh, sizeof(log), nullptr, log);
            (void)log; // could log via ImGui
            glDeleteShader(sh);
            return 0;
        }
        return sh;
    };
    GLuint v = compile(GL_VERTEX_SHADER, kVert);
    if (!v) return 0;
    GLuint f = compile(GL_FRAGMENT_SHADER, kFrag);
    if (!f) { glDeleteShader(v); return 0; }
    GLuint p = glCreateProgram();
    glAttachShader(p, v);
    glAttachShader(p, f);
    glLinkProgram(p);
    glDeleteShader(v);
    glDeleteShader(f);
    GLint ok = 0;
    glGetProgramiv(p, GL_LINK_STATUS, &ok);
    if (!ok) { glDeleteProgram(p); return 0; }
    return p;
}

void ensure_init(Cache& c) {
    if (c.initialized) return;
    fn::gfx::gl_loader_init();
    fn::gfx::fb_init(c.fb);
    c.program = compile_program();
    if (c.program) {
        c.loc_view  = glGetUniformLocation(c.program, "u_view");
        c.loc_proj  = glGetUniformLocation(c.program, "u_proj");
        c.loc_color = glGetUniformLocation(c.program, "u_color");
    }
    c.initialized = true;
}

} // namespace

void mesh_viewer(const char* id, const MeshViewerConfig& cfg) {
    if (!id) id = "##mesh_viewer";

    // Resolve panel size.
    ImVec2 size = cfg.size;
    if (size.x <= 0) size.x = ImGui::GetContentRegionAvail().x;
    if (size.y <= 0) size.y = 400.0f;
    int w = (int)size.x; if (w < 1) w = 1;
    int h = (int)size.y; if (h < 1) h = 1;

    auto& c = caches()[id];
    ensure_init(c);

    if (cfg.mesh && cfg.mesh->ok() && cfg.cam && c.program) {
        cfg.cam->aspect = (float)w / (float)h;

        fn::gfx::fb_resize(c.fb, w, h);

        // Save GL state.
        GLint prev_fbo = 0;
        glGetIntegerv(GL_FRAMEBUFFER_BINDING, &prev_fbo);
        GLint prev_vp[4];
        glGetIntegerv(GL_VIEWPORT, prev_vp);
        GLboolean prev_depth = glIsEnabled(GL_DEPTH_TEST);

        glBindFramebuffer(GL_FRAMEBUFFER, c.fb.fbo);
        glViewport(0, 0, w, h);
        glClearColor(0.10f, 0.10f, 0.13f, 1.0f);
        glClear(GL_COLOR_BUFFER_BIT);
        // No depth attachment in our FBO — fall back to back-to-front-ish via
        // GL_DEPTH_TEST off. For inspection meshes this is fine; documented.
        glDisable(GL_DEPTH_TEST);

        glUseProgram(c.program);
        auto m = fn::core::orbit_camera_matrices(*cfg.cam);
        // Row-major -> upload with transpose=GL_TRUE.
        if (c.loc_view >= 0) glUniformMatrix4fv(c.loc_view, 1, GL_TRUE, m.view);
        if (c.loc_proj >= 0) glUniformMatrix4fv(c.loc_proj, 1, GL_TRUE, m.proj);
        if (c.loc_color >= 0) {
            float r = ((cfg.color >>  0) & 0xFF) / 255.0f;
            float g = ((cfg.color >>  8) & 0xFF) / 255.0f;
            float b = ((cfg.color >> 16) & 0xFF) / 255.0f;
            float a = ((cfg.color >> 24) & 0xFF) / 255.0f;
            glUniform4f(c.loc_color, r, g, b, a);
        }

#ifndef __EMSCRIPTEN__
        if (cfg.wireframe) {
            glPolygonMode(GL_FRONT_AND_BACK, GL_LINE);
        } else {
            glPolygonMode(GL_FRONT_AND_BACK, GL_FILL);
        }
#endif

        glBindVertexArray(cfg.mesh->vao);
        glDrawElements(GL_TRIANGLES, cfg.mesh->index_count, GL_UNSIGNED_INT, 0);
        glBindVertexArray(0);

#ifndef __EMSCRIPTEN__
        glPolygonMode(GL_FRONT_AND_BACK, GL_FILL);
#endif
        glUseProgram(0);

        // Restore GL state.
        glBindFramebuffer(GL_FRAMEBUFFER, (GLuint)prev_fbo);
        glViewport(prev_vp[0], prev_vp[1], prev_vp[2], prev_vp[3]);
        if (prev_depth) glEnable(GL_DEPTH_TEST);
    }

    // Display.
    ImGui::Image((ImTextureID)(intptr_t)c.fb.tex,
                 ImVec2((float)w, (float)h),
                 ImVec2(0, 1), ImVec2(1, 0));

    if (ImGui::IsItemActive() && cfg.cam) {
        ImVec2 d = ImGui::GetIO().MouseDelta;
        fn::core::orbit_camera_handle_drag(*cfg.cam, d, 0.0f);
    }
    if (ImGui::IsItemHovered() && cfg.cam) {
        float w_scroll = ImGui::GetIO().MouseWheel;
        if (w_scroll != 0.0f) {
            ImVec2 zero{0, 0};
            fn::core::orbit_camera_handle_drag(*cfg.cam, zero, w_scroll);
        }
    }
}

} // namespace fn::viz