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

    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() {
        return this.free.pop();
    }

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

export class AtlasAllocator {
    atlases = [];

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

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

    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();
            if (allocation != null) {
                return { i, allocation };
            }
        }

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

        return { i, allocation };
    }

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

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

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

        // 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[id.i].download(gl, id);
        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;
                }
            }
        }
    }
}