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;

/**
 * 重构后的 ModelRender：更模块化、健壮的渲染子系统
 */
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 = GL20.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\n" +
                    "layout(location = 0) in vec2 aPosition;\n" +
                    "layout(location = 1) in vec2 aTexCoord;\n" +
                    "out vec2 vTexCoord;\n" +
                    "uniform mat3 uModelMatrix;\n" +
                    "uniform mat3 uViewMatrix;\n" +
                    "uniform mat3 uProjectionMatrix;\n" +
                    "void main() {\n" +
                    "    vec3 p = uProjectionMatrix * uViewMatrix * uModelMatrix * vec3(aPosition, 1.0);\n" +
                    "    gl_Position = vec4(p.xy, 0.0, 1.0);\n" +
                    "    vTexCoord = aTexCoord;\n" +
                    "}";

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

    // ================== 初始化 / 清理 ==================
    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();

        // 设置投影与视图（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 的 GL 资源已上传（ModelRender 管理 upload）
        MeshGLResources res = meshResources.computeIfAbsent(mesh, k -> new MeshGLResources());
        if (!res.initialized) uploadMeshData(mesh, res);

        // 绑定纹理到单元0（我们使用 0 固定）
        Texture tex = mesh.getTexture();
        int texId = (tex != null && !tex.isDisposed()) ? tex.getTextureId() : defaultTextureId;

        // active unit & bind — 确保 shader 已被 use()（调用者保证）
        GL13.glActiveTexture(GL13.GL_TEXTURE0);
        GL11.glBindTexture(GL11.GL_TEXTURE_2D, texId);
        // 将 sampler 设为 0（内部函数保证 program 绑定）
        setUniformIntInternal(defaultProgram, "uTexture", 0);

        // 绑定 VAO 并绘制
        GL30.glBindVertexArray(res.vao);
        int drawMode = getGLDrawMode(mesh.getDrawMode());
        if (mesh.getIndices().length > 0 &&
                (drawMode == GL11.GL_TRIANGLES || drawMode == GL11.GL_TRIANGLE_STRIP || drawMode == GL11.GL_TRIANGLE_FAN)) {
            GL11.glDrawElements(drawMode, mesh.getIndices().length, GL11.GL_UNSIGNED_INT, 0);
        } else {
            GL11.glDrawArrays(drawMode, 0, res.vertexCount);
        }
        GL30.glBindVertexArray(0);

        // 解绑纹理（避免污染后续 state）
        GL11.glBindTexture(GL11.GL_TEXTURE_2D, 0);

        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);
        for (int i = 0; i < vertexCount; i++) {
            inter.put(verts[i*2]);
            inter.put(verts[i*2+1]);
            inter.put(uvs[i*2]);
            inter.put(uvs[i*2+1]);
        }
        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(); }

}