import {mat4Identity, mat4Multiply, mat4Ortho, mat4LookAt, mat4Translate} from './math.js';
import {loadBlockstate, loadModel, TextureAtlas} from './assets.js';
import {resolveTexture, buildGeometry} from './geometry.js';
import {resolveBlock, getVariantHash} from './blockstate.js';

const VS_SRC = `#version 300 es
layout(location=0) in vec3 a_position;
layout(location=1) in vec3 a_normal;
layout(location=2) in vec2 a_uv;
layout(location=3) in float a_tint;
layout(location=4) in float a_shade; // 0.0 = emissive, 1.0 = shaded

// Instanced attributes shifted to account for a_shade
layout(location=5) in mat4 i_matrix;
layout(location=9) in vec3 i_color;

uniform mat4 u_viewProj;
uniform vec3 u_lightDir;

out vec2 v_uv;
out float v_light;
out vec3 v_color;

void main() {
    gl_Position = u_viewProj * i_matrix * vec4(a_position, 1.0);
    v_uv = a_uv;

    vec3 normal = normalize(mat3(i_matrix) * a_normal);
    
    // Lambertian lighting
    float baseLight = max(dot(normal, normalize(u_lightDir)), 0.0) * 0.6 + 0.4;
    
    // Mix between full brightness (1.0) and shaded based on the element's shade flag
    v_light = mix(1.0, baseLight, a_shade);
    
    v_color = mix(vec3(1.0), i_color, a_tint);
}
`;

const FS_SRC = `#version 300 es
precision highp float;
in vec2 v_uv;
in float v_light;
in vec3 v_color;
uniform sampler2D u_texture;
out vec4 fragColor;
void main() {
    vec4 texColor = texture(u_texture, v_uv);
    if (texColor.a < 0.1) discard;
    fragColor = vec4(texColor.rgb * v_color * v_light, texColor.a);
}
`;

function compileShader(gl, type, src) {
    const shader = gl.createShader(type);
    gl.shaderSource(shader, src);
    gl.compileShader(shader);
    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {
        throw new Error(gl.getShaderInfoLog(shader));
    }
    return shader;
}

export class Renderer {
    constructor(canvas) {
        this.gl = canvas.getContext('webgl2', {antialias: true});
        if (!this.gl) throw new Error("WebGL2 not supported");

        const gl = this.gl;
        gl.enable(gl.DEPTH_TEST);
        gl.enable(gl.CULL_FACE);

        const vs = compileShader(gl, gl.VERTEX_SHADER, VS_SRC);
        const fs = compileShader(gl, gl.FRAGMENT_SHADER, FS_SRC);
        this.program = gl.createProgram();
        gl.attachShader(this.program, vs);
        gl.attachShader(this.program, fs);
        gl.linkProgram(this.program);

        this.uniforms = {
            viewProj: gl.getUniformLocation(this.program, "u_viewProj"),
            lightDir: gl.getUniformLocation(this.program, "u_lightDir"),
            texture: gl.getUniformLocation(this.program, "u_texture")
        };

        this.atlasTexture = gl.createTexture();
        this.atlas = new TextureAtlas(512);
        this.meshes = []; // Array of renderable VAO objects

        // Default Camera Setup
        this.projMatrix = new Float32Array(16);
        this.viewMatrix = new Float32Array(16);
        this.viewProjMatrix = new Float32Array(16);
        this.setCamera(-4, 4, -4, 4, 8, 6, 8, 1, 0, 1);
    }

    setCamera(left, right, bottom, top, eyeX, eyeY, eyeZ, targetX, targetY, targetZ) {
        mat4Ortho(this.projMatrix, left, right, bottom, top, -20, 20);
        mat4LookAt(this.viewMatrix, eyeX, eyeY, eyeZ, targetX, targetY, targetZ, 0, 1, 0);
        mat4Multiply(this.viewProjMatrix, this.projMatrix, this.viewMatrix);
    }

    // This method takes a World object and rebuilds the WebGL buffers
    async updateWorld(world) {
        const instancePool = new Map();

        // 1. Resolve all blocks in the world into part definitions
        for (const block of world.blocks.values()) {
            const stateJSON = await loadBlockstate(block.id);
            const parts = resolveBlock(stateJSON, block.props);

            for (const part of parts) {
                const hash = getVariantHash(part);
                if (!instancePool.has(hash)) {
                    instancePool.set(hash, {partDef: part, matrices: [], colors: []});
                }

                const matrix = mat4Identity(new Float32Array(16));
                mat4Translate(matrix, matrix, [block.x, block.y, block.z]);

                const pool = instancePool.get(hash);
                pool.matrices.push(...matrix);

                if (block.props.power) {
                    const p = parseInt(block.props.power, 10);
                    pool.colors.push((0x4B + (p * 12)) / 255, 0.0, 0.0);
                } else {
                    pool.colors.push(1.0, 1.0, 1.0);
                }
            }
        }

        // 2. Build geometries, load textures, and create WebGL buffers
        const newMeshes = [];
        for (const [hash, pool] of instancePool.entries()) {
            const modelJSON = await loadModel(pool.partDef.model);

            for (const el of modelJSON.elements || []) {
                for (const face of Object.values(el.faces || {})) {
                    const texPath = resolveTexture(modelJSON, face.texture);
                    if (texPath) await this.atlas.load(texPath);
                }
            }

            const geometry = buildGeometry(modelJSON, this.atlas, pool.partDef);
            const matrixArray = new Float32Array(pool.matrices);
            const colorArray = new Float32Array(pool.colors);
            const instanceCount = pool.matrices.length / 16;

            newMeshes.push(this.createInstancedMesh(geometry, instanceCount, matrixArray, colorArray));
        }

        // 3. Upload updated atlas and swap the active meshes
        this.updateAtlasTexture();
        this.meshes = newMeshes;
    }

    updateAtlasTexture() {
        const gl = this.gl;
        gl.bindTexture(gl.TEXTURE_2D, this.atlasTexture);
        gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, this.atlas.canvas);
        gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
        gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
    }

    createInstancedMesh(geometry, instanceCount, matrices, colors) {
        const gl = this.gl;
        const vao = gl.createVertexArray();
        gl.bindVertexArray(vao);

        // 1. Bind standard geometry (divisor = 0)
        const bindGeomAttr = (loc, data, size) => {
            const buffer = gl.createBuffer();
            gl.bindBuffer(gl.ARRAY_BUFFER, buffer);
            gl.bufferData(gl.ARRAY_BUFFER, data, gl.STATIC_DRAW);
            gl.enableVertexAttribArray(loc);
            gl.vertexAttribPointer(loc, size, gl.FLOAT, false, 0, 0);
        };

        bindGeomAttr(0, geometry.positions, 3);
        bindGeomAttr(1, geometry.normals, 3);
        bindGeomAttr(2, geometry.uvs, 2);
        bindGeomAttr(3, geometry.tints, 1);
        bindGeomAttr(4, geometry.shades, 1);

        // 2. Bind Index Buffer
        const ebo = gl.createBuffer();
        gl.bindBuffer(gl.ELEMENT_ARRAY_BUFFER, ebo);
        gl.bufferData(gl.ELEMENT_ARRAY_BUFFER, geometry.indices, gl.STATIC_DRAW);

        // 3. Bind Instanced Attributes (divisor = 1)
        const matrixBuffer = gl.createBuffer();
        gl.bindBuffer(gl.ARRAY_BUFFER, matrixBuffer);
        gl.bufferData(gl.ARRAY_BUFFER, matrices, gl.STATIC_DRAW);

        // Mat4 requires 4 separate vec4 attributes in WebGL
        for (let i = 0; i < 4; i++) {
            const loc = 5 + i;
            gl.enableVertexAttribArray(loc);
            gl.vertexAttribPointer(loc, 4, gl.FLOAT, false, 64, i * 16);
            gl.vertexAttribDivisor(loc, 1);
        }

        const colorBuffer = gl.createBuffer();
        gl.bindBuffer(gl.ARRAY_BUFFER, colorBuffer);
        gl.bufferData(gl.ARRAY_BUFFER, colors, gl.STATIC_DRAW);
        gl.enableVertexAttribArray(9);
        gl.vertexAttribPointer(9, 3, gl.FLOAT, false, 0, 0);
        gl.vertexAttribDivisor(9, 1);

        gl.bindVertexArray(null);

        return {vao, indexCount: geometry.indices.length, instanceCount};
    }

    render() {
        const gl = this.gl;
        gl.viewport(0, 0, gl.canvas.width, gl.canvas.height);
        gl.clearColor(0.1, 0.1, 0.12, 1.0);
        gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);

        gl.useProgram(this.program);
        gl.uniformMatrix4fv(this.uniforms.viewProj, false, this.viewProjMatrix);
        gl.uniform3f(this.uniforms.lightDir, 1.0, 2.0, 0.5);

        gl.activeTexture(gl.TEXTURE0);
        gl.bindTexture(gl.TEXTURE_2D, this.atlasTexture);
        gl.uniform1i(this.uniforms.texture, 0);

        for (const mesh of this.meshes) {
            gl.bindVertexArray(mesh.vao);
            gl.drawElementsInstanced(gl.TRIANGLES, mesh.indexCount, gl.UNSIGNED_SHORT, 0, mesh.instanceCount);
        }
    }
}