package net.minecraft.client.renderer;

import com.mojang.blaze3d.buffers.GpuBufferSlice;
import com.mojang.blaze3d.systems.GpuDevice;
import com.mojang.blaze3d.systems.RenderSystem;
import com.mojang.logging.LogUtils;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import net.minecraft.util.Mth;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;

public class DynamicUniformStorage<T extends DynamicUniformStorage.DynamicUniform> implements AutoCloseable {
	private static final Logger LOGGER = LogUtils.getLogger();
	private final List<MappableRingBuffer> oldBuffers = new ArrayList();
	private final int blockSize;
	private MappableRingBuffer ringBuffer;
	private int nextBlock;
	private int capacity;
	@Nullable
	private T lastUniform;
	private final String label;

	public DynamicUniformStorage(final String label, final int uboSize, final int initialCapacity) {
		GpuDevice device = RenderSystem.getDevice();
		this.blockSize = Mth.roundToward(uboSize, device.getDeviceInfo().limits().minUniformOffsetAlignment());
		this.capacity = Mth.smallestEncompassingPowerOfTwo(initialCapacity);
		this.nextBlock = 0;
		this.ringBuffer = new MappableRingBuffer(() -> label + " x" + this.blockSize, 130, this.blockSize * this.capacity);
		this.label = label;
	}

	public void endFrame() {
		this.nextBlock = 0;
		this.lastUniform = null;
		this.ringBuffer.rotate();
		if (!this.oldBuffers.isEmpty()) {
			for (MappableRingBuffer oldBuffer : this.oldBuffers) {
				oldBuffer.close();
			}

			this.oldBuffers.clear();
		}
	}

	private void resizeBuffers(final int newCapacity) {
		this.capacity = newCapacity;
		this.nextBlock = 0;
		this.lastUniform = null;
		this.oldBuffers.add(this.ringBuffer);
		this.ringBuffer = new MappableRingBuffer(() -> this.label + " x" + this.blockSize, 130, this.blockSize * this.capacity);
	}

	public GpuBufferSlice writeUniform(final T uniform) {
		if (this.lastUniform != null && this.lastUniform.equals(uniform)) {
			return this.ringBuffer.currentBuffer().slice((this.nextBlock - 1) * this.blockSize, this.blockSize);
		} else {
			if (this.nextBlock >= this.capacity) {
				int newCapacity = this.capacity * 2;
				LOGGER.info("Resizing {}, capacity limit of {} reached during a single frame. New capacity will be {}.", this.label, this.capacity, newCapacity);
				this.resizeBuffers(newCapacity);
			}

			int offset = this.nextBlock * this.blockSize;

			try (GpuBufferSlice.MappedView view = this.ringBuffer.currentBuffer().slice(offset, this.blockSize).map(false, true)) {
				uniform.write(view.data());
			}

			this.nextBlock++;
			this.lastUniform = uniform;
			return this.ringBuffer.currentBuffer().slice(offset, this.blockSize);
		}
	}

	public GpuBufferSlice[] writeUniforms(final T[] uniforms) {
		if (uniforms.length == 0) {
			return new GpuBufferSlice[0];
		} else {
			if (this.nextBlock + uniforms.length > this.capacity) {
				int newCapacity = Mth.smallestEncompassingPowerOfTwo(Math.max(this.capacity + 1, uniforms.length));
				LOGGER.info("Resizing {}, capacity limit of {} reached during a single frame. New capacity will be {}.", this.label, this.capacity, newCapacity);
				this.resizeBuffers(newCapacity);
			}

			int firstOffset = this.nextBlock * this.blockSize;
			GpuBufferSlice[] result = new GpuBufferSlice[uniforms.length];

			try (GpuBufferSlice.MappedView view = this.ringBuffer.currentBuffer().slice(firstOffset, uniforms.length * this.blockSize).map(false, true)) {
				ByteBuffer byteBuffer = view.data();

				for (int i = 0; i < uniforms.length; i++) {
					T uniform = uniforms[i];
					result[i] = this.ringBuffer.currentBuffer().slice(firstOffset + i * this.blockSize, this.blockSize);
					byteBuffer.position(i * this.blockSize);
					uniform.write(byteBuffer);
				}
			}

			this.nextBlock += uniforms.length;
			this.lastUniform = uniforms[uniforms.length - 1];
			return result;
		}
	}

	public void close() {
		for (MappableRingBuffer oldBuffer : this.oldBuffers) {
			oldBuffer.close();
		}

		this.ringBuffer.close();
	}

	public interface DynamicUniform {
		void write(ByteBuffer byteBuffer);
	}
}