// engine.js
import {mat4Identity, mat4Translate} from './math.js';
import {Cache} from "./cache.js";
import {Blockstate, BlockstateHandler} from "./blockstate.js";
import {ModelHandler} from "./model.js";
import {TextureHandler} from "./texture.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;

layout(location=5) in mat4 i_matrix;
layout(location=9) in vec3 i_color;

uniform mat4 u_viewProj;
uniform vec3 u_lightDir;
uniform vec3 u_viewDir;

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);
    
    float sun = max(dot(normal, normalize(u_lightDir)), 0.0);
    float head = max(dot(normal, normalize(u_viewDir)), 0.0);
    float baseLight = clamp(sun + head * 0.3, 0.0, 1.0) * 0.6 + 0.4;
    
    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 Engine {
    constructor() {
        this.canvas = document.createElement('canvas');
        this.canvas.style.position = 'fixed';
        this.canvas.style.top = '0';
        this.canvas.style.left = '0';
        this.canvas.style.width = '100vw';
        this.canvas.style.height = '100vh';
        this.canvas.style.zIndex = '-1';
        document.body.appendChild(this.canvas);

        this.gl = this.canvas.getContext('webgl2', {antialias: true, alpha: true});
        const gl = this.gl;
        gl.enable(gl.DEPTH_TEST);
        gl.enable(gl.CULL_FACE);
        gl.enable(gl.SCISSOR_TEST);

        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"),
            viewDir: gl.getUniformLocation(this.program, "u_viewDir"),
            texture: gl.getUniformLocation(this.program, "u_texture")
        };

        this.atlasTexture = gl.createTexture();

        // --- Initialize the new Resource Pipeline ---
        this.cache = new Cache('assets');
        this.atlas = new TextureHandler(256);
        this.cache.register('blockstates', new BlockstateHandler());
        this.cache.register('models', new ModelHandler());
        this.cache.register('textures', this.atlas);

        this.dioramas = [];
        this.worldMeshes = new Map();
        this.renderRequested = false;

        this.observedWorlds = new Set();
        this.updateRequested = false;

        window.addEventListener('resize', () => {
            this.resize();
            this.requestRender();
        });
        window.addEventListener('scroll', () => this.requestRender(), {passive: true});
        this.resize();
    }

    requestUpdate() {
        if (!this.updateRequested) {
            this.updateRequested = true;
            // Add 'async' here and 'await' updateAll
            Promise.resolve().then(async () => {
                await this.updateAll();
                this.updateRequested = false;
            });
        }
    }

    requestRender() {
        if (!this.renderRequested) {
            this.renderRequested = true;
            requestAnimationFrame(() => {
                this.renderRequested = false;
                this.render();
            });
        }
    }

    resize() {
        this.canvas.width = window.innerWidth;
        this.canvas.height = window.innerHeight;
    }

    addDiorama(diorama) {
        this.dioramas.push(diorama);
        if (!this.observedWorlds.has(diorama.worlds)) {
            this.observedWorlds.add(diorama.world);
            diorama.world.subscribe(() => this.requestUpdate());
        }
    }

    async updateAll() {
        const uniqueWorlds = new Set(this.dioramas.map(d => d.world));
        if (uniqueWorlds.size === 0) return;

        // --- SWEEP 1: Collect & Fetch Blockstates ---
        const uniqueBlockIds = new Set();
        for (const world of uniqueWorlds) {
            for (const block of world.blocks.values()) uniqueBlockIds.add(block.id);
        }
        await Promise.all(Array.from(uniqueBlockIds).map(id => this.cache.get(id, 'blockstates')));

        // --- SWEEP 2: Resolve Parts & Fetch Models ---
        const uniqueModelIds = new Set();
        const parsedWorlds = new Map();

        for (const world of uniqueWorlds) {
            const parsedBlocks = [];
            for (const block of world.blocks.values()) {
                const state = await this.cache.get(block.id, 'blockstates');
                const parts = state.resolveParts(block.state);
                parsedBlocks.push({block, parts});
                for (const p of parts) uniqueModelIds.add(p.model);
            }
            parsedWorlds.set(world, parsedBlocks);
        }
        await Promise.all(Array.from(uniqueModelIds).map(id => this.cache.get(id, 'models')));

        // --- SWEEP 3: Pool Instances & Fetch Textures ---
        const textureTasks = [];
        const worldPools = new Map();

        for (const world of uniqueWorlds) {
            const instancePool = new Map();
            for (const {block, parts} of parsedWorlds.get(world)) {
                for (const part of parts) {
                    const hash = Blockstate.getVariantHash(part);
                    if (!instancePool.has(hash)) {
                        instancePool.set(hash, {partDef: part, matrices: [], colors: []});
                    }

                    const matrix = mat4Identity(new Float32Array(16));
                    mat4Translate(matrix, matrix, block.pos);
                    const pool = instancePool.get(hash);
                    pool.matrices.push(...matrix);

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

            for (const pool of instancePool.values()) {
                const blockModel = await this.cache.get(pool.partDef.model, 'models');
                for (const el of blockModel.elements) {
                    for (const face of Object.values(el.faces || {})) {
                        const texPath = blockModel.resolveTexture(face.texture);
                        if (texPath) {
                            // Synchronously allocates UV, asynchronously fetches image
                            textureTasks.push(this.cache.get(texPath, 'textures', 'png'));
                        }
                    }
                }
            }
            worldPools.set(world, instancePool);
        }

        await Promise.all(textureTasks);
        this.updateAtlasTexture();

        // --- SWEEP 4: Bake Geometry ---
        for (const world of uniqueWorlds) {
            const newMeshes = [];
            for (const pool of worldPools.get(world).values()) {
                const blockModel = await this.cache.get(pool.partDef.model, 'models');
                const geometry = blockModel.buildGeometry(this.atlas.uvmap, 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));
            }
            this.worldMeshes.set(world, newMeshes);
        }

        this.requestRender();
    }

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

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

        const ebo = gl.createBuffer();
        gl.bindBuffer(gl.ELEMENT_ARRAY_BUFFER, ebo);
        gl.bufferData(gl.ELEMENT_ARRAY_BUFFER, geometry.indices, gl.STATIC_DRAW);

        const matrixBuffer = gl.createBuffer();
        gl.bindBuffer(gl.ARRAY_BUFFER, matrixBuffer);
        gl.bufferData(gl.ARRAY_BUFFER, matrices, gl.STATIC_DRAW);

        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.scissor(0, 0, gl.canvas.width, gl.canvas.height);
        gl.clearColor(0.0, 0.0, 0.0, 0.0);
        gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);

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

        for (const diorama of this.dioramas) {
            const rect = diorama.element.getBoundingClientRect();
            if (rect.bottom < 0 || rect.top > gl.canvas.height || rect.right < 0 || rect.left > gl.canvas.width) continue;

            const bottom = gl.canvas.height - rect.bottom;
            gl.viewport(rect.left, bottom, rect.width, rect.height);
            gl.scissor(rect.left, bottom, rect.width, rect.height);

            const aspect = rect.width / rect.height;
            const viewProj = diorama.updateMatrices(aspect);

            gl.uniformMatrix4fv(this.uniforms.viewProj, false, viewProj);
            gl.uniform3f(this.uniforms.lightDir, 1.0, 3.0, 2.0);
            gl.uniform3f(this.uniforms.viewDir, diorama.viewDir[0], diorama.viewDir[1], diorama.viewDir[2]);

            const meshes = this.worldMeshes.get(diorama.world) || [];
            for (const mesh of meshes) {
                gl.bindVertexArray(mesh.vao);
                gl.drawElementsInstanced(gl.TRIANGLES, mesh.indexCount, gl.UNSIGNED_SHORT, 0, mesh.instanceCount);
            }
        }
    }
}