window-axis-innovators-box/src/main/java/com/chuangzhou/vivid2D/render/ModelRender.java

package com.chuangzhou.vivid2D.render;

import com.chuangzhou.vivid2D.render.model.Model2D;
import com.chuangzhou.vivid2D.render.model.ModelPart;
import com.chuangzhou.vivid2D.render.model.util.Mesh2D;
import com.chuangzhou.vivid2D.render.model.util.Texture;
import org.joml.Matrix3f;
import org.joml.Vector2f;
import org.joml.Vector4f;
import org.lwjgl.opengl.*;
import org.lwjgl.system.MemoryUtil;

import java.nio.ByteBuffer;
import java.nio.FloatBuffer;
import java.nio.IntBuffer;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

import static org.lwjgl.opengl.GL20.glGetUniformLocation;

/**
 * 重构后的 ModelRender：更模块化、健壮的渲染子系统
 * @author tzdwindows 7
 */
public final class ModelRender {

    private ModelRender() { /* no instances */ }

    // ================== 全局状态 ==================
    private static boolean initialized = false;
    private static int viewportWidth = 800;
    private static int viewportHeight = 600;
    private static final Vector4f CLEAR_COLOR = new Vector4f(0.0f, 0.0f, 0.0f, 1.0f);
    private static boolean enableDepthTest = false;
    private static boolean enableBlending = true;

    // 着色器与资源
    private static final Map<String, ShaderProgram> shaderMap = new HashMap<>();
    private static ShaderProgram defaultProgram = null;

    private static final Map<Mesh2D, MeshGLResources> meshResources = new HashMap<>();
    private static final AtomicInteger textureUnitAllocator = new AtomicInteger(0);

    // 默认白色纹理
    private static int defaultTextureId = 0;

    // ================== 内部类：ShaderProgram ==================
    private static class ShaderProgram {
        final int programId;
        final Map<String, Integer> uniformCache = new HashMap<>();

        ShaderProgram(int programId) {
            this.programId = programId;
        }

        void use() {
            GL20.glUseProgram(programId);
        }

        void stop() {
            GL20.glUseProgram(0);
        }

        int getUniformLocation(String name) {
            return uniformCache.computeIfAbsent(name, k -> {
                int loc = glGetUniformLocation(programId, k);
                if (loc == -1) {
                    // debug 时可以打开
                    // System.err.println("Warning: uniform not found: " + k);
                }
                return loc;
            });
        }

        void delete() {
            if (GL20.glIsProgram(programId)) GL20.glDeleteProgram(programId);
        }
    }

    // ================== 内部类：MeshGLResources ==================
    private static class MeshGLResources {
        int vao = 0;
        int vbo = 0;
        int ebo = 0;
        int vertexCount = 0;
        boolean initialized = false;

        void dispose() {
            if (ebo != 0) { GL15.glDeleteBuffers(ebo); ebo = 0; }
            if (vbo != 0) { GL15.glDeleteBuffers(vbo); vbo = 0; }
            if (vao != 0) { GL30.glDeleteVertexArrays(vao); vao = 0; }
            initialized = false;
        }
    }

    // ================== 着色器源 ==================
    private static final String VERTEX_SHADER_SRC =
            """
                    #version 330 core
                    layout(location = 0) in vec2 aPosition;
                    layout(location = 1) in vec2 aTexCoord;
                    out vec2 vTexCoord;
                    out vec2 vDebugPos;
                    uniform mat3 uModelMatrix;
                    uniform mat3 uViewMatrix;
                    uniform mat3 uProjectionMatrix;
                    void main() {
                        vec3 p = uProjectionMatrix * uViewMatrix * uModelMatrix * vec3(aPosition, 1.0);
                        gl_Position = vec4(p.xy, 0.0, 1.0);
                        vTexCoord = aTexCoord;
                        vDebugPos = p.xy;
                    }""";

    private static final String FRAGMENT_SHADER_SRC =
            """
                    #version 330 core
                    in vec2 vTexCoord;
                    in vec2 vDebugPos;
                    out vec4 FragColor;
                    uniform sampler2D uTexture;
                    uniform vec4 uColor;
                    uniform float uOpacity;
                    uniform int uBlendMode;
                    uniform int uDebugMode;
                    void main() {
                        if (uDebugMode == 1) {
                            FragColor = vec4(vDebugPos * 0.5 + 0.5, 0.0, 1.0);
                            return;
                        }
                        
                        vec4 tex = texture(uTexture, vTexCoord);
                        vec4 finalColor = tex * uColor;
                        if (uBlendMode == 1) finalColor.rgb = tex.rgb + uColor.rgb;
                        else if (uBlendMode == 2) finalColor.rgb = tex.rgb * uColor.rgb;
                        else if (uBlendMode == 3) finalColor.rgb = 1.0 - (1.0 - tex.rgb) * (1.0 - uColor.rgb);
                        finalColor.a = tex.a * uOpacity;
                        if (finalColor.a <= 0.001) discard;
                        FragColor = finalColor;
                    }""";

    // ================== 初始化 / 清理 ==================
    public static synchronized void initialize() {
        if (initialized) return;

        System.out.println("Initializing ModelRender...");

        // 需要在外部创建 OpenGL 上下文并调用 GL.createCapabilities()
        logGLInfo();

        // 初始 GL 状态
        setupGLState();

        // 创建默认 shader
        try {
            compileDefaultShader();
        } catch (RuntimeException ex) {
            System.err.println("Failed to compile default shader: " + ex.getMessage());
            throw ex;
        }

        // 创建默认纹理
        createDefaultTexture();

        // 初始化视口
        GL11.glViewport(0, 0, viewportWidth, viewportHeight);

        initialized = true;
        System.out.println("ModelRender initialized successfully");
    }

    private static void logGLInfo() {
        System.out.println("OpenGL Vendor: " + GL11.glGetString(GL11.GL_VENDOR));
        System.out.println("OpenGL Renderer: " + GL11.glGetString(GL11.GL_RENDERER));
        System.out.println("OpenGL Version: " + GL11.glGetString(GL11.GL_VERSION));
        System.out.println("GLSL Version: " + GL20.glGetString(GL20.GL_SHADING_LANGUAGE_VERSION));
    }

    private static void setupGLState() {
        GL11.glClearColor(CLEAR_COLOR.x, CLEAR_COLOR.y, CLEAR_COLOR.z, CLEAR_COLOR.w);

        if (enableBlending) {
            GL11.glEnable(GL11.GL_BLEND);
            GL11.glBlendFunc(GL11.GL_SRC_ALPHA, GL11.GL_ONE_MINUS_SRC_ALPHA);
        } else {
            GL11.glDisable(GL11.GL_BLEND);
        }

        if (enableDepthTest) {
            GL11.glEnable(GL11.GL_DEPTH_TEST);
            GL11.glDepthFunc(GL11.GL_LEQUAL);
        } else {
            GL11.glDisable(GL11.GL_DEPTH_TEST);
        }

        GL11.glDisable(GL11.GL_CULL_FACE);

        checkGLError("setupGLState");
    }

    private static void compileDefaultShader() {
        int vs = compileShader(GL20.GL_VERTEX_SHADER, VERTEX_SHADER_SRC);
        int fs = compileShader(GL20.GL_FRAGMENT_SHADER, FRAGMENT_SHADER_SRC);
        int prog = linkProgram(vs, fs);
        ShaderProgram sp = new ShaderProgram(prog);
        shaderMap.put("default", sp);
        defaultProgram = sp;

        // 设置一些默认 uniform（需要先 use）
        sp.use();
        setUniformIntInternal(sp, "uTexture", 0);
        setUniformFloatInternal(sp, "uOpacity", 1.0f);
        setUniformVec4Internal(sp, "uColor", new Vector4f(1,1,1,1));
        setUniformIntInternal(sp, "uBlendMode", 0);
        sp.stop();
    }

    private static int compileShader(int type, String src) {
        int shader = GL20.glCreateShader(type);
        GL20.glShaderSource(shader, src);
        GL20.glCompileShader(shader);
        int status = GL20.glGetShaderi(shader, GL20.GL_COMPILE_STATUS);
        if (status == GL11.GL_FALSE) {
            String log = GL20.glGetShaderInfoLog(shader);
            GL20.glDeleteShader(shader);
            throw new RuntimeException("Shader compilation failed: " + log);
        }
        return shader;
    }

    private static int linkProgram(int vs, int fs) {
        int prog = GL20.glCreateProgram();
        GL20.glAttachShader(prog, vs);
        GL20.glAttachShader(prog, fs);
        GL20.glLinkProgram(prog);
        int status = GL20.glGetProgrami(prog, GL20.GL_LINK_STATUS);
        if (status == GL11.GL_FALSE) {
            String log = GL20.glGetProgramInfoLog(prog);
            GL20.glDeleteProgram(prog);
            throw new RuntimeException("Program link failed: " + log);
        }
        // shaders can be deleted after linking
        GL20.glDetachShader(prog, vs);
        GL20.glDetachShader(prog, fs);
        GL20.glDeleteShader(vs);
        GL20.glDeleteShader(fs);
        return prog;
    }

    private static void createDefaultTexture() {
        // 使用 GL11.glGenTextures() 获取单个 id（更直观，避免 IntBuffer 问题）
        defaultTextureId = GL11.glGenTextures();
        GL11.glBindTexture(GL11.GL_TEXTURE_2D, defaultTextureId);
        ByteBuffer white = MemoryUtil.memAlloc(4);
        white.put((byte)255).put((byte)255).put((byte)255).put((byte)255).flip();
        GL11.glTexImage2D(GL11.GL_TEXTURE_2D, 0, GL11.GL_RGBA8, 1, 1, 0, GL11.GL_RGBA, GL11.GL_UNSIGNED_BYTE, white);
        MemoryUtil.memFree(white);
        GL11.glTexParameteri(GL11.GL_TEXTURE_2D, GL11.GL_TEXTURE_MIN_FILTER, GL11.GL_NEAREST);
        GL11.glTexParameteri(GL11.GL_TEXTURE_2D, GL11.GL_TEXTURE_MAG_FILTER, GL11.GL_NEAREST);
        GL11.glBindTexture(GL11.GL_TEXTURE_2D, 0);

        checkGLError("createDefaultTexture");
    }

    public static synchronized void cleanup() {
        if (!initialized) return;

        System.out.println("Cleaning up ModelRender...");

        // mesh resources
        for (MeshGLResources r : meshResources.values()) r.dispose();
        meshResources.clear();

        // shaders
        for (ShaderProgram sp : shaderMap.values()) sp.delete();
        shaderMap.clear();
        defaultProgram = null;

        // textures
        if (defaultTextureId != 0) {
            GL11.glDeleteTextures(defaultTextureId);
            defaultTextureId = 0;
        }

        initialized = false;
        System.out.println("ModelRender cleaned up");
    }

    // ================== 渲染流程 ==================
    public static void render(float deltaTime, Model2D model) {
        if (!initialized) throw new IllegalStateException("ModelRender not initialized");
        if (model == null) return;

        // 更新模型（确保 worldTransform 已经被计算）
        model.update(deltaTime);

        GL11.glClear(GL11.GL_COLOR_BUFFER_BIT | (enableDepthTest ? GL11.GL_DEPTH_BUFFER_BIT : 0));

        // 使用默认 shader（保持绑定直到完成渲染）
        defaultProgram.use();

        // setUniformIntInternal(defaultProgram, "uDebugMode", 0); 设置debug模式

        // 设置投影与视图（3x3 正交投影用于 2D）
        Matrix3f proj = buildOrthoProjection(viewportWidth, viewportHeight);
        setUniformMatrix3(defaultProgram, "uProjectionMatrix", proj);
        setUniformMatrix3(defaultProgram, "uViewMatrix", new Matrix3f().identity());

        // 递归渲染所有根部件（使用 3x3 矩阵）
        Matrix3f identity = new Matrix3f().identity();
        for (ModelPart p : model.getParts()) {
            if (p.getParent() != null) continue;
            renderPartRecursive(p, identity);
        }

        defaultProgram.stop();

        checkGLError("render");
    }

    /**
     * 关键修改点：在渲染前确保更新 part 的 worldTransform，
     * 然后直接使用 part.getWorldTransform() 作为 uModelMatrix 传入 shader。
     */
    private static void renderPartRecursive(ModelPart part, Matrix3f parentMat) {
        // 确保 part 的 local/world 矩阵被计算（会更新 transformDirty）
        part.updateWorldTransform(parentMat, false);

        // 直接使用已经计算好的 worldTransform
        Matrix3f world = part.getWorldTransform();

        // 先设置部件相关的 uniform（opacity / blendMode / color 等）
        setPartUniforms(defaultProgram, part);

        // 把 world 矩阵传给 shader（兼容 uModelMatrix 和 可能的 uModel）
        setUniformMatrix3(defaultProgram, "uModelMatrix", world);
        setUniformMatrix3(defaultProgram, "uModel", world);

        // 绘制本节点的所有 mesh（将 world 传入 renderMesh）
        for (Mesh2D mesh : part.getMeshes()) {
            renderMesh(mesh, world);
        }

        // 递归渲染子节点，继续传入当前 world 作为子节点的 parent
        for (ModelPart child : part.getChildren()) {
            renderPartRecursive(child, world);
        }
    }

    private static void renderMesh(Mesh2D mesh, Matrix3f modelMatrix) {
        // 使用默认 shader（保证 shader 已被 use）
        defaultProgram.use();

        // 如果 mesh 已经被烘焙到世界坐标，则传 identity 矩阵给 shader（防止重复变换）
        Matrix3f matToUse = mesh.isBakedToWorld() ? new Matrix3f().identity() : modelMatrix;

        // 确保 shader 中的矩阵 uniform 已更新（再次设置以防遗漏）
        setUniformMatrix3(defaultProgram, "uModelMatrix", matToUse);
        setUniformMatrix3(defaultProgram, "uModel", matToUse);

        // 调用 Mesh2D 的 draw 重载（传 program id 与实际矩阵）
        try {
            mesh.draw(defaultProgram.programId, matToUse);
        } catch (AbstractMethodError | NoSuchMethodError e) {
            // 回退：仍然兼容旧的无参 draw（在这种情况下 shader 的 uModelMatrix 已经被设置）
            mesh.draw();
        }

        checkGLError("renderMesh");
    }


    private static int getGLDrawMode(int meshDrawMode) {
        switch (meshDrawMode) {
            case Mesh2D.POINTS: return GL11.GL_POINTS;
            case Mesh2D.LINES: return GL11.GL_LINES;
            case Mesh2D.LINE_STRIP: return GL11.GL_LINE_STRIP;
            case Mesh2D.TRIANGLES: return GL11.GL_TRIANGLES;
            case Mesh2D.TRIANGLE_STRIP: return GL11.GL_TRIANGLE_STRIP;
            case Mesh2D.TRIANGLE_FAN: return GL11.GL_TRIANGLE_FAN;
            default: return GL11.GL_TRIANGLES;
        }
    }

    // ================== 上传数据 ==================（被弃用）


    // ================== uniform 设置辅助（内部使用，确保 program 已绑定） ==================
    private static void setUniformIntInternal(ShaderProgram sp, String name, int value) {
        int loc = sp.getUniformLocation(name);
        if (loc != -1) GL20.glUniform1i(loc, value);
    }

    private static void setUniformFloatInternal(ShaderProgram sp, String name, float value) {
        int loc = sp.getUniformLocation(name);
        if (loc != -1) GL20.glUniform1f(loc, value);
    }

    private static void setUniformVec4Internal(ShaderProgram sp, String name, Vector4f vec) {
        int loc = sp.getUniformLocation(name);
        if (loc != -1) GL20.glUniform4f(loc, vec.x, vec.y, vec.z, vec.w);
    }

    private static void setUniformMatrix3(ShaderProgram sp, String name, Matrix3f m) {
        int loc = sp.getUniformLocation(name);
        if (loc == -1) return;
        FloatBuffer fb = MemoryUtil.memAllocFloat(9);
        try {
            m.get(fb);
            GL20.glUniformMatrix3fv(loc, false, fb);
        } finally {
            MemoryUtil.memFree(fb);
        }
    }


    // 外部可用的统一设置（会自动切换到默认程序）
    private static void setUniformInt(String name, int value) {
        defaultProgram.use();
        setUniformIntInternal(defaultProgram, name, value);
        defaultProgram.stop();
    }
    private static void setUniformFloat(String name, float value) {
        defaultProgram.use();
        setUniformFloatInternal(defaultProgram, name, value);
        defaultProgram.stop();
    }
    private static void setUniformVec4(String name, Vector4f v) {
        defaultProgram.use();
        setUniformVec4Internal(defaultProgram, name, v);
        defaultProgram.stop();
    }

    // ================== 部件属性 ==================
    private static void setPartUniforms(ShaderProgram sp, ModelPart part) {
        setUniformFloatInternal(sp, "uOpacity", part.getOpacity());
        int blend = 0;
        ModelPart.BlendMode bm = part.getBlendMode();
        if (bm != null) {
            switch (bm) {
                case ADDITIVE: blend = 1; break;
                case MULTIPLY: blend = 2; break;
                case SCREEN: blend = 3; break;
                case NORMAL: default: blend = 0;
            }
        } else {
            blend = 0;
        }
        setUniformIntInternal(sp, "uBlendMode", blend);
        // 这里保留为白色，若需要部件 tint 请替换为 part 的 color 属性
        setUniformVec4Internal(sp, "uColor", new Vector4f(1,1,1,1));
    }

    // ================== 工具 ==================
    private static Matrix3f buildOrthoProjection(int width, int height) {
        Matrix3f m = new Matrix3f();
        m.set(
                2.0f / width, 0.0f, -1.0f,
                0.0f, -2.0f / height, 1.0f,
                0.0f, 0.0f, 1.0f
        );
        return m;
    }

    public static void setViewport(int width, int height) {
        viewportWidth = Math.max(1, width);
        viewportHeight = Math.max(1, height);
        GL11.glViewport(0, 0, viewportWidth, viewportHeight);
    }

    public static void setClearColor(float r, float g, float b, float a) {
        GL11.glClearColor(r,g,b,a);
    }

    private static void checkGLError(String op) {
        int e = GL11.glGetError();
        if (e != GL11.GL_NO_ERROR) {
            System.err.println("OpenGL error during " + op + ": " + getGLErrorString(e));
        }
    }

    private static String getGLErrorString(int err) {
        switch (err) {
            case GL11.GL_INVALID_ENUM: return "GL_INVALID_ENUM";
            case GL11.GL_INVALID_VALUE: return "GL_INVALID_VALUE";
            case GL11.GL_INVALID_OPERATION: return "GL_INVALID_OPERATION";
            case GL11.GL_OUT_OF_MEMORY: return "GL_OUT_OF_MEMORY";
            default: return "Unknown(0x" + Integer.toHexString(err) + ")";
        }
    }

    // ================== 辅助：外部获取状态 ==================
    public static boolean isInitialized() { return initialized; }
    public static int getLoadedMeshCount() { return meshResources.size(); }

}