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;
uniform vec3 u_upDir;

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;
    
    // Strict Schematic Detection: Both vectors must be perfectly axis-aligned
    vec3 absV = abs(u_viewDir);
    vec3 absU = abs(u_upDir);
    bool isSchematic = (absV.x + absV.y + absV.z == 1.0) && (absU.x + absU.y + absU.z == 1.0);
    
    if (isSchematic && head == 1.0) {
        baseLight = 1.0;
    }
    
    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.display = 'none';
        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);

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

        this.atlasTexture = gl.createTexture();

        this.cache = new Cache('assets', () => {
            for (const frame of this.observedFrames) this.dirtyFrames.add(frame);
            for (const diorama of this.dioramas) diorama.dirty = true;
            this.requestUpdate();
        });
        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.frameMeshes = new Map();

        this.renderRequested = false;
        this.updateRequested = false;

        this.observedFrames = new Set();
        this.dirtyFrames = new Set(); // Tracks specific timelines that need mesh updates

        window.addEventListener('resize', () => this.requestRender());
        window.addEventListener('scroll', () => this.requestRender());
    }

    requestUpdate() {
        if (!this.updateRequested) {
            this.updateRequested = true;
            Promise.resolve().then(async () => {
                this.updateAll();
                this.updateRequested = false;
            });
        }
    }

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

    addDiorama(diorama) {
        this.dioramas.push(diorama);
        diorama.dirty = true;

        if (!this.observedFrames.has(diorama.frame)) {
            this.observedFrames.add(diorama.frame);
            this.dirtyFrames.add(diorama.frame);

            this.preloadWorld(diorama.frame.world);

            diorama.frame.subscribe(() => {
                this.dirtyFrames.add(diorama.frame);
                for (const d of this.dioramas) {
                    if (d.frame === diorama.frame) d.dirty = true;
                }
                this.requestUpdate();
            });
        } else {
            this.requestRender();
        }
    }

    async preloadWorld(world) {
        const blocks = world.getUniqueBlockStates();

        // 1. Await all Blockstates
        const uniqueIds = new Set(blocks.map(b => b.id));
        await Promise.all(Array.from(uniqueIds).map(id => this.cache.getAsync(id, 'blockstates')));

        // 2. Resolve permutations and await all Models
        const uniqueModels = new Set();
        for (const block of blocks) {
            const stateDef = this.cache.getSync(block.id, 'blockstates');
            const parts = stateDef.resolveParts(block.state);
            for (const p of parts) uniqueModels.add(p.model);
        }
        await Promise.all(Array.from(uniqueModels).map(id => this.cache.getAsync(id, 'models')));

        // 3. Scan geometry and await all Textures
        const textureTasks = [];
        for (const modelId of uniqueModels) {
            const blockModel = this.cache.getSync(modelId, 'models');
            for (const el of blockModel.elements) {
                for (const face of Object.values(el.faces || {})) {
                    const texPath = blockModel.resolveTexture(face.texture);
                    if (texPath && texPath !== ':missing') {
                        textureTasks.push(this.cache.getAsync(texPath, 'textures'));
                    }
                }
            }
        }
        await Promise.all(textureTasks);
    }

    updateAll() {
        if (this.dirtyFrames.size === 0) return;

        const framesToUpdate = Array.from(this.dirtyFrames);
        this.dirtyFrames.clear();

        const parsedFrames = new Map();

        for (const frame of framesToUpdate) {
            const parsedBlocks = [];
            for (const block of frame.blocks.values()) {
                // Sync grab (returns fallback instantly if loading)
                const state = this.cache.getSync(block.id, 'blockstates');
                const parts = state.resolveParts(block.state);
                parsedBlocks.push({block, parts});
            }
            parsedFrames.set(frame, parsedBlocks);
        }

        const framePools = new Map();

        for (const frame of framesToUpdate) {
            const instancePool = new Map();
            for (const {block, parts} of parsedFrames.get(frame)) {
                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 = this.cache.getSync(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) {
                            this.cache.getSync(texPath, 'textures');
                        }
                    }
                }
            }
            framePools.set(frame, instancePool);
        }

        this.updateAtlasTexture();

        for (const frame of framesToUpdate) {
            const oldMeshes = this.frameMeshes.get(frame);
            if (oldMeshes) {
                for (const mesh of oldMeshes) {
                    this.gl.deleteVertexArray(mesh.vao);
                    for (const buf of mesh.buffers) this.gl.deleteBuffer(buf);
                }
            }

            const newMeshes = [];
            for (const pool of framePools.get(frame).values()) {
                const blockModel = this.cache.getSync(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.frameMeshes.set(frame, newMeshes);
        }

        for (const d of this.dioramas) {
            if (framesToUpdate.includes(d.frame)) d.dirty = true;
        }

        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 buffers = []; // Track buffers for GC

        const bindGeomAttr = (loc, data, size) => {
            const buffer = gl.createBuffer();
            buffers.push(buffer);
            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();
        buffers.push(ebo);
        gl.bindBuffer(gl.ELEMENT_ARRAY_BUFFER, ebo);
        gl.bufferData(gl.ELEMENT_ARRAY_BUFFER, geometry.indices, gl.STATIC_DRAW);

        const matrixBuffer = gl.createBuffer();
        buffers.push(matrixBuffer);
        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();
        buffers.push(colorBuffer);
        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, buffers, indexCount: geometry.indices.length, instanceCount};
    }

    render() {
        const gl = this.gl;
        let setupProgram = false;

        for (const diorama of this.dioramas) {
            const rect = diorama.canvas.getBoundingClientRect();

            // Visibility Culling
            if (rect.bottom < 0 || rect.top > window.innerHeight || rect.right < 0 || rect.left > window.innerWidth) {
                continue;
            }

            const dpr = window.devicePixelRatio || 1;
            const targetW = Math.floor(rect.width * dpr);
            const targetH = Math.floor(rect.height * dpr);

            if (diorama.canvas.width !== targetW || diorama.canvas.height !== targetH) {
                diorama.canvas.width = targetW;
                diorama.canvas.height = targetH;
                diorama.dirty = true;
            }

            // Lazy Render: Only draw if the diorama actually changed
            if (!diorama.dirty) continue;

            if (!setupProgram) {
                gl.useProgram(this.program);
                gl.activeTexture(gl.TEXTURE0);
                gl.bindTexture(gl.TEXTURE_2D, this.atlasTexture);
                gl.uniform1i(this.uniforms.texture, 0);
                setupProgram = true;
            }

            if (this.canvas.width < targetW || this.canvas.height < targetH) {
                this.canvas.width = Math.max(this.canvas.width, targetW);
                this.canvas.height = Math.max(this.canvas.height, targetH);
            }

            gl.viewport(0, 0, targetW, targetH);
            gl.clearColor(0.0, 0.0, 0.0, 0.0);
            gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);

            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]);
            gl.uniform3f(this.uniforms.upDir, diorama.upDir[0], diorama.upDir[1], diorama.upDir[2]);

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

            diorama.ctx2d.clearRect(0, 0, targetW, targetH);
            diorama.ctx2d.drawImage(
                this.canvas,
                0, 0, targetW, targetH,
                0, 0, targetW, targetH
            );

            // Clean state!
            diorama.dirty = false;
        }
    }
}