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");
    }

    private static void renderPartRecursive(ModelPart part, Matrix3f parentMat) {
        Matrix3f local = part.getLocalTransform();
        Matrix3f world = new Matrix3f(parentMat).mul(local);  // world = parent * local

        // 从 world 矩阵取世界坐标
        //float worldX = world.m02;
        //float worldY = world.m12;
        //System.out.println("Rendering part: " + part.getName() + " at world position: " + worldX + ", " + worldY);
        // 传入 shader
        setUniformMatrix3(defaultProgram, "uModelMatrix", world);
        setPartUniforms(defaultProgram, part);

        for (Mesh2D mesh : part.getMeshes()) {
            renderMesh(mesh);
        }

        for (ModelPart child : part.getChildren()) {
            renderPartRecursive(child, world);
        }
    }


    private static void renderMesh(Mesh2D mesh) {
        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;
        }
    }

    // ================== 上传数据 ==================（被弃用）
    //private static void uploadMeshData(Mesh2D mesh, MeshGLResources res) {
    //    System.out.println("Uploading mesh data: " + mesh.getName());
//
    //    res.vao = GL30.glGenVertexArrays();
    //    GL30.glBindVertexArray(res.vao);
//
    //    res.vbo = GL15.glGenBuffers();
    //    GL15.glBindBuffer(GL15.GL_ARRAY_BUFFER, res.vbo);
//
    //    float[] verts = mesh.getVertices();
    //    float[] uvs = mesh.getUVs();
    //    int vertexCount = mesh.getVertexCount();
    //    if (verts == null || verts.length == 0) throw new IllegalStateException("Mesh has no vertices: " + mesh.getName());
//
    //    FloatBuffer inter = MemoryUtil.memAllocFloat(vertexCount * 4);
//
    //    // ========== 添加调试输出 ==========
    //    System.out.println("=== Vertex data debug output ===");
    //    System.out.println("Grid name: " + mesh.getName());
    //    System.out.println("Number of vertices: " + vertexCount);
    //    System.out.println("Vertex coordinates (x, y):");
//
    //    for (int i = 0; i < vertexCount; i++) {
    //        float x = verts[i*2];
    //        float y = verts[i*2+1];
    //        float u = uvs[i*2];
    //        float v = uvs[i*2+1];
//
    //        // 输出每个顶点的坐标和UV
    //        System.out.printf("vertex %d: location(%.3f, %.3f), UV(%.3f, %.3f)%n",
    //                i, x, y, u, v);
//
    //        inter.put(x);
    //        inter.put(y);
    //        inter.put(u);
    //        inter.put(v);
    //    }
    //    System.out.println("=== Vertex data output is over ===");
    //    // ========== 调试输出结束 ==========
//
    //    inter.flip();
    //    GL15.glBufferData(GL15.GL_ARRAY_BUFFER, inter, GL15.GL_STATIC_DRAW);
    //    MemoryUtil.memFree(inter);
//
    //    // 设置 attribute（位置 / uv），layout 已在 shader 中固定
    //    int stride = 4 * Float.BYTES;
    //    GL20.glEnableVertexAttribArray(0);
    //    GL20.glVertexAttribPointer(0, 2, GL11.GL_FLOAT, false, stride, 0);
    //    GL20.glEnableVertexAttribArray(1);
    //    GL20.glVertexAttribPointer(1, 2, GL11.GL_FLOAT, false, stride, 2 * Float.BYTES);
//
    //    int[] indices = mesh.getIndices();
    //    if (indices != null && indices.length > 0) {
    //        res.ebo = GL15.glGenBuffers();
    //        GL15.glBindBuffer(GL15.GL_ELEMENT_ARRAY_BUFFER, res.ebo);
    //        IntBuffer ib = MemoryUtil.memAllocInt(indices.length);
    //        ib.put(indices).flip();
    //        GL15.glBufferData(GL15.GL_ELEMENT_ARRAY_BUFFER, ib, GL15.GL_STATIC_DRAW);
    //        MemoryUtil.memFree(ib);
    //        res.vertexCount = indices.length; // drawElements 使用 count
    //    } else {
    //        res.vertexCount = vertexCount;
    //    }
//
    //    // 不解绑 ELEMENT_ARRAY_BUFFER（它属于 VAO），解绑 ARRAY_BUFFER
    //    GL15.glBindBuffer(GL15.GL_ARRAY_BUFFER, 0);
    //    GL30.glBindVertexArray(0);
//
    //    res.initialized = true;
    //    checkGLError("uploadMeshData");
    //    System.out.println("Uploaded mesh: " + mesh.getName() + " (v=" + vertexCount + ")");
    //}

    // ================== 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;
        switch (part.getBlendMode()) {
            case ADDITIVE: blend = 1; break;
            case MULTIPLY: blend = 2; break;
            case SCREEN: blend = 3; break;
            case NORMAL: default: 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(); }

}