import { compileProgram } from "rkgk/webgl.js";

class Atlas {
    static getInitBuffer(chunkSize, nChunks) {
        let imageSize = chunkSize * nChunks;
        return new Uint8Array(imageSize * imageSize * 4).fill(0x00);
    }

    constructor(gl, chunkSize, nChunks, initBuffer) {
        this.chunkSize = chunkSize;
        this.nChunks = nChunks;
        this.textureSize = chunkSize * nChunks;

        this.free = Array(nChunks * nChunks);
        for (let y = 0; y < nChunks; ++y) {
            for (let x = 0; x < nChunks; ++x) {
                this.free[x + y * nChunks] = { x, y };
            }
        }

        this.texture = gl.createTexture();
        gl.bindTexture(gl.TEXTURE_2D, this.texture);
        gl.texImage2D(
            gl.TEXTURE_2D,
            0,
            gl.RGBA8,
            this.textureSize,
            this.textureSize,
            0,
            gl.RGBA,
            gl.UNSIGNED_BYTE,
            initBuffer,
        );
        gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
        gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);

        this.framebuffer = gl.createFramebuffer();
        gl.bindFramebuffer(gl.FRAMEBUFFER, this.framebuffer);
        gl.framebufferTexture2D(
            gl.FRAMEBUFFER,
            gl.COLOR_ATTACHMENT0,
            gl.TEXTURE_2D,
            this.texture,
            0,
        );
        gl.bindFramebuffer(gl.FRAMEBUFFER, null);
    }

    alloc(gl, initBuffer) {
        let xy = this.free.pop();
        if (xy != null) {
            this.upload(gl, xy, initBuffer);
        }
        return xy;
    }

    dealloc(xy) {
        this.free.push(xy);
    }

    upload(gl, { x, y }, source) {
        gl.bindTexture(gl.TEXTURE_2D, this.texture);
        gl.texSubImage2D(
            gl.TEXTURE_2D,
            0,
            x * this.chunkSize,
            y * this.chunkSize,
            this.chunkSize,
            this.chunkSize,
            gl.RGBA,
            gl.UNSIGNED_BYTE,
            source,
        );
    }

    // Assumes a pack PBO is bound.
    download(gl, { x, y }) {
        gl.bindFramebuffer(gl.FRAMEBUFFER, this.framebuffer);
        gl.readPixels(
            x * this.chunkSize,
            y * this.chunkSize,
            this.chunkSize,
            this.chunkSize,
            gl.RGBA,
            gl.UNSIGNED_BYTE,
            0,
        );
        gl.bindFramebuffer(gl.FRAMEBUFFER, null);
    }

    getFramebufferRect({ x, y }) {
        return {
            x: x * this.chunkSize,
            y: y * this.chunkSize,
            width: this.chunkSize,
            height: this.chunkSize,
        };
    }
}

const compositeVertexShader = `#version 300 es
    precision highp float;

    layout (location = 0) in vec2 a_position;
    layout (location = 1) in vec2 a_uv;

    out vec2 vf_uv;

    void main() {
        gl_Position = vec4(a_position, 0.0, 1.0);
        vf_uv = a_uv;
    }
`;

const compositeFragmentShader = `#version 300 es
    precision highp float;

    uniform sampler2D u_chunk;

    in vec2 vf_uv;

    out vec4 f_color;

    void main() {
        f_color = texture(u_chunk, vf_uv);
        // f_color = vec4(vec3(0.0), 1.0);
    }
`;

export class AtlasAllocator {
    atlases = [];

    // Allocation names
    #ids = new Map();
    #idCounter = 1;

    // Download buffers
    #pboPool = [];
    #downloadBufferPool = [];
    #pendingDownloads = [];

    constructor(gl, chunkSize, nChunks) {
        this.gl = gl;
        this.chunkSize = chunkSize;
        this.nChunks = nChunks;
        this.initBuffer = Atlas.getInitBuffer(chunkSize, nChunks);

        // Compositing pipeline

        let compositeProgramId = compileProgram(gl, compositeVertexShader, compositeFragmentShader);
        this.compositeProgram = {
            id: compositeProgramId,
            u_chunk: gl.getUniformLocation(compositeProgramId, "u_chunk"),
        };

        // prettier-ignore
        this.compositeRectData = new Float32Array([
            // a_position
            -1, 1, // 0: top left
            1, 1, // 1: top right
            1, -1, // 2: bottom right
            -1, -1, // 3: bottom left

            // a_uv - filled out later when compositing
            0, 0,
            0, 0,
            0, 0,
            0, 0,
        ]);
        let compositeRectIndices = new Uint16Array([0, 1, 2, 2, 3, 0]);
        this.compositeRectUv = this.compositeRectData.subarray(8);

        this.compositeRectVao = gl.createVertexArray();
        this.compositeRectVbo = gl.createBuffer();
        this.compositeRectIbo = gl.createBuffer();
        gl.bindVertexArray(this.compositeRectVao);

        gl.bindBuffer(gl.ARRAY_BUFFER, this.compositeRectVbo);
        gl.bufferData(gl.ARRAY_BUFFER, this.compositeRectData, gl.DYNAMIC_DRAW);
        console.log(this.compositeRectData.byteLength);

        gl.bindBuffer(gl.ELEMENT_ARRAY_BUFFER, this.compositeRectIbo);
        gl.bufferData(gl.ELEMENT_ARRAY_BUFFER, compositeRectIndices, gl.DYNAMIC_DRAW);

        gl.vertexAttribPointer(0, 2, gl.FLOAT, false, 2 * 4, 0);
        gl.vertexAttribPointer(1, 2, gl.FLOAT, false, 2 * 4, this.compositeRectUv.byteOffset);
        for (let i = 0; i < 2; ++i) gl.enableVertexAttribArray(i);

        this.compositeTexture = gl.createTexture();
        gl.bindTexture(gl.TEXTURE_2D, this.compositeTexture);
        gl.texImage2D(
            gl.TEXTURE_2D,
            0,
            gl.RGBA8,
            chunkSize,
            chunkSize,
            0,
            gl.RGBA,
            gl.UNSIGNED_BYTE,
            null,
        );
        gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
        gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);

        this.compositeFramebuffer = gl.createFramebuffer();
        gl.bindFramebuffer(gl.FRAMEBUFFER, this.compositeFramebuffer);
        gl.framebufferTexture2D(
            gl.FRAMEBUFFER,
            gl.COLOR_ATTACHMENT0,
            gl.TEXTURE_2D,
            this.compositeTexture,
            0,
        );
        gl.bindFramebuffer(gl.FRAMEBUFFER, null);
    }

    #obtainId(allocInfo) {
        let id = this.#idCounter++;
        this.#ids.set(id, allocInfo);
        return id;
    }

    #releaseId(id) {
        this.#ids.delete(id);
    }

    #getAllocInfo(id) {
        return this.#ids.get(id);
    }

    getAtlasIndex(id) {
        return this.#getAllocInfo(id).i;
    }

    getAllocation(id) {
        return this.#getAllocInfo(id).allocation;
    }

    alloc() {
        // Right now we do a dumb linear scan through all atlases, but in the future it would be
        // really nice to optimize this by storing information about which atlases have free slots
        // precisely.

        for (let i = 0; i < this.atlases.length; ++i) {
            let atlas = this.atlases[i];
            let allocation = atlas.alloc(this.gl, this.initBuffer);
            if (allocation != null) {
                return this.#obtainId({ i, allocation });
            }
        }

        let i = this.atlases.length;
        let atlas = new Atlas(this.gl, this.chunkSize, this.nChunks, this.initBuffer);
        let allocation = atlas.alloc(this.gl, this.initBuffer);
        this.atlases.push(atlas);

        return this.#obtainId({ i, allocation });
    }

    dealloc(id) {
        let { i, allocation } = this.#getAllocInfo(id);
        let atlas = this.atlases[i];
        atlas.dealloc(allocation);
        this.#releaseId(id);
    }

    upload(id, source) {
        let { i, allocation } = this.#getAllocInfo(id);
        this.atlases[i].upload(this.gl, allocation, source);
    }

    async download(id) {
        let gl = this.gl;
        let allocInfo = this.#getAllocInfo(id);

        // Get PBO

        let pbo = this.#pboPool.pop();
        if (pbo == null) {
            let dataSize = this.chunkSize * this.chunkSize * 4;
            pbo = gl.createBuffer();
            gl.bindBuffer(gl.PIXEL_PACK_BUFFER, pbo);
            gl.bufferData(gl.PIXEL_PACK_BUFFER, dataSize, gl.DYNAMIC_READ);
        }

        // Initiate download

        gl.bindBuffer(gl.PIXEL_PACK_BUFFER, pbo);
        this.atlases[allocInfo.i].download(gl, allocInfo.allocation);
        let fence = gl.fenceSync(gl.SYNC_GPU_COMMANDS_COMPLETE, 0);

        gl.bindBuffer(gl.PIXEL_PACK_BUFFER, null);

        // Add for ticking

        return new Promise((resolve) => {
            this.#pendingDownloads.push({ pbo, fence, resolve });
        });
    }

    // Call after download() finishes running to give memory back to the allocator, for reuse in
    // later pixel transfers.
    freeDownload(arrayBuffer) {
        this.#downloadBufferPool.push(arrayBuffer);
    }

    // Call every frame to poll for download completion.
    tickDownloads() {
        if (this.#pendingDownloads.length == 0) return;

        let gl = this.gl;

        for (let i = 0; i < this.#pendingDownloads.length; ++i) {
            let pending = this.#pendingDownloads[i];
            let status = gl.getSyncParameter(pending.fence, gl.SYNC_STATUS);
            if (status == gl.SIGNALED) {
                // Transfer complete, fetch pixels back to an array buffer.
                let dataSize = this.chunkSize * this.chunkSize * 4;
                let arrayBuffer = this.#downloadBufferPool.pop() ?? new ArrayBuffer(dataSize);

                gl.bindBuffer(gl.PIXEL_PACK_BUFFER, pending.pbo);
                gl.getBufferSubData(gl.PIXEL_PACK_BUFFER, 0, new Uint8Array(arrayBuffer));
                gl.bindBuffer(gl.PIXEL_PACK_BUFFER, null);
                gl.deleteSync(pending.fence);

                pending.resolve({
                    width: this.chunkSize,
                    height: this.chunkSize,
                    data: arrayBuffer,
                });

                let last = this.#pendingDownloads.pop();
                if (this.#pendingDownloads.length > 0) {
                    this.#pendingDownloads[i] = last;
                    --i;
                } else {
                    break; // now empty
                }
            }
        }
    }

    canvasSource() {
        let useCanvas = (gl, id) => {
            let allocInfo = this.#getAllocInfo(id);
            let atlas = this.atlases[allocInfo.i];

            let viewport = atlas.getFramebufferRect(allocInfo.allocation);

            gl.enable(gl.SCISSOR_TEST);
            gl.bindFramebuffer(gl.FRAMEBUFFER, atlas.framebuffer);
            gl.viewport(viewport.x, viewport.y, viewport.width, viewport.height);
            gl.scissor(viewport.x, viewport.y, viewport.width, viewport.height);

            return viewport;
        };

        let resetCanvas = (gl) => {
            gl.bindFramebuffer(gl.FRAMEBUFFER, null);
            gl.disable(gl.SCISSOR_TEST);
        };

        return {
            useCanvas,
            resetCanvas,
        };
    }

    // NOTE: I was thinking a bit about whether the chunk allocator is the right place to put
    // compositing operations like this. After much consideration, I've decided that it's pretty
    // much the only sensible place, because it's the only place concerned with the layout of
    // chunks in memory, and rendering of laid out chunks is quite implementation dependent.
    //
    // This does break the purity of the "allocator" role a bit though, but I don't really know if
    // there's a good way around that.
    //
    // Maybe. But I don't feel like breaking this apart to 10 smaller classes, either.

    #drawComposite(u, v, uvScale) {
        // Assumes bound source texture, destination framebuffer, and viewport set.

        let gl = this.gl;

        gl.bindVertexArray(this.compositeRectVao);
        gl.bindBuffer(gl.ARRAY_BUFFER, this.compositeRectVbo);
        gl.useProgram(this.compositeProgram.id);

        gl.uniform1i(this.compositeProgram.u_chunk, 0);

        let uv = this.compositeRectUv;
        uv[0] = u * uvScale;
        uv[1] = v * uvScale;
        uv[2] = (u + 1) * uvScale;
        uv[3] = v * uvScale;
        uv[4] = (u + 1) * uvScale;
        uv[5] = (v + 1) * uvScale;
        uv[6] = u * uvScale;
        uv[7] = (v + 1) * uvScale;
        gl.bufferSubData(gl.ARRAY_BUFFER, uv.byteOffset, uv);

        gl.drawElements(gl.TRIANGLES, 6, gl.UNSIGNED_SHORT, 0);
    }

    composite(dstId, srcId, op) {
        // NOTE: This has to go through an intermediate buffer in case the source and destination
        // atlas are the same.

        console.assert(op == "alphaBlend", "composite operation must be alphaBlend");

        let gl = this.gl;

        let dstAtlas = this.atlases[this.getAtlasIndex(dstId)];
        let srcAtlas = this.atlases[this.getAtlasIndex(srcId)];
        let dstAllocation = this.getAllocation(dstId);
        let srcAllocation = this.getAllocation(srcId);

        // Source -> intermediate buffer

        gl.disable(gl.BLEND);
        gl.disable(gl.SCISSOR_TEST);

        gl.bindFramebuffer(gl.FRAMEBUFFER, this.compositeFramebuffer);
        gl.bindTexture(gl.TEXTURE_2D, srcAtlas.texture);
        gl.viewport(0, 0, this.chunkSize, this.chunkSize);
        this.#drawComposite(srcAllocation.x, srcAllocation.y, 1 / this.nChunks);

        // Intermediate buffer -> destination

        gl.enable(gl.BLEND);
        gl.blendFunc(gl.SRC_ALPHA, gl.ONE_MINUS_SRC_ALPHA);

        gl.bindFramebuffer(gl.FRAMEBUFFER, dstAtlas.framebuffer);
        gl.bindTexture(gl.TEXTURE_2D, this.compositeTexture);
        gl.viewport(
            dstAllocation.x * this.chunkSize,
            dstAllocation.y * this.chunkSize,
            this.chunkSize,
            this.chunkSize,
        );
        this.#drawComposite(0, 0, 1);

        // Cleanup

        gl.bindFramebuffer(gl.FRAMEBUFFER, null);
    }
}