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,
            null,
        );
        gl.bindFramebuffer(gl.FRAMEBUFFER, null);
    }

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

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

    #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, 0);

        // 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.
    freeDownloaded(arrayBuffer) {
        this.#downloadBufferPool.push(arrayBuffer);
    }

    // Call every frame to poll for download completion.
    tickDownloads() {
        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, arrayBuffer);
                gl.bindBuffer(gl.PIXEL_PACK_BUFFER, 0);
                gl.deleteSync(pending.fence);

                pending.resolve(arrayBuffer);

                let last = this.#pendingDownloads.pop();
                if (last != null) {
                    this.#pendingDownloads[i] = last;
                    --i;
                }
            }
        }
    }

    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,
        };
    }
}