package com.mojang.blaze3d.vulkan;

import com.mojang.blaze3d.buffers.GpuBuffer;
import com.mojang.blaze3d.buffers.GpuBufferSlice.MappedView;
import java.nio.ByteBuffer;
import java.nio.LongBuffer;
import java.util.Objects;
import java.util.function.Supplier;
import org.jspecify.annotations.Nullable;
import org.lwjgl.PointerBuffer;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.system.MemoryUtil;
import org.lwjgl.util.vma.Vma;
import org.lwjgl.util.vma.VmaAllocationCreateInfo;
import org.lwjgl.vulkan.VkBufferCreateInfo;

public class VulkanGpuBuffer extends GpuBuffer implements Destroyable {
	private final VulkanDevice device;
	private boolean closed = false;
	private final long vkBuffer;
	private final long vmaAllocation;
	private int mappingRefCount = 0;

	public VulkanGpuBuffer(
		final VulkanDevice device,
		@Nullable final Supplier<String> label,
		@GpuBuffer.Usage final int usage,
		final long size,
		final boolean forceHostVisibleAllocation
	) {
		super(usage, size);
		this.device = device;

		try (MemoryStack stack = MemoryStack.stackPush()) {
			VkBufferCreateInfo bufferCreateInfo = VkBufferCreateInfo.calloc(stack).sType$Default();
			bufferCreateInfo.size(size);
			bufferCreateInfo.usage(VulkanConst.bufferUsageToVk(usage));
			bufferCreateInfo.sharingMode(0);
			bufferCreateInfo.pQueueFamilyIndices(null);
			VmaAllocationCreateInfo allocCreateInfo = VmaAllocationCreateInfo.calloc(stack);
			allocCreateInfo.usage(8);
			if (forceHostVisibleAllocation) {
				allocCreateInfo.requiredFlags(allocCreateInfo.requiredFlags() | 2);
			}

			if (VulkanUtils.hasAnyBit(usage, 3)) {
				allocCreateInfo.requiredFlags(allocCreateInfo.requiredFlags() | 2 | 4);
				if (VulkanUtils.hasAnyBit(usage, 1)) {
					allocCreateInfo.preferredFlags(allocCreateInfo.preferredFlags() | 8);
					allocCreateInfo.flags(allocCreateInfo.flags() | 2048);
				} else {
					allocCreateInfo.flags(allocCreateInfo.flags() | 1024);
				}
			}

			if (VulkanUtils.hasAnyBit(usage, 5)) {
				allocCreateInfo.usage(9);
			}

			LongBuffer bufferPtr = stack.callocLong(1);
			PointerBuffer allocPtr = stack.callocPointer(1);
			int result = Vma.vmaCreateBuffer(device.vma(), bufferCreateInfo, allocCreateInfo, bufferPtr, allocPtr, null);
			VulkanUtils.crashIfFailure(device, result, "Failed to allocate VkBuffer");
			this.vkBuffer = bufferPtr.get(0);
			this.vmaAllocation = allocPtr.get(0);
			if (label != null) {
				device.instance().debug().setObjectName(device.vkDevice(), 9, this.vkBuffer, label);
			}
		}
	}

	@Override
	public void destroy() {
		Vma.vmaDestroyBuffer(this.device.vma(), this.vkBuffer, this.vmaAllocation);
	}

	@Override
	public boolean isClosed() {
		return this.closed;
	}

	@Override
	public void close() {
		if (!this.closed) {
			this.closed = true;
			if (this.mappingRefCount != 0) {
				throw new IllegalStateException("Attempt to close a mapped buffer");
			} else {
				this.device.createCommandEncoder().queueForDestroy(this);
			}
		}
	}

	public long vkBuffer() {
		return this.vkBuffer;
	}

	@Override
	public MappedView map(final long offset, final long length, final boolean read, final boolean write) {
		if (this.isClosed()) {
			throw new IllegalStateException("Buffer already closed");
		} else if (!read && !write) {
			throw new IllegalArgumentException("At least read or write must be true");
		} else if (read && (this.usage() & 1) == 0) {
			throw new IllegalStateException("Buffer is not readable");
		} else if (write && (this.usage() & 2) == 0) {
			throw new IllegalStateException("Buffer is not writable");
		} else if (offset + length > this.size()) {
			throw new IllegalArgumentException(
				"Cannot map more data than this buffer can hold (attempting to map " + length + " bytes at offset " + offset + " from " + this.size() + " size buffer)"
			);
		} else if (length > 2147483647L) {
			throw new IllegalArgumentException("Mapping buffer slice larger than 2GB is not supported");
		} else if (offset >= 0L && length >= 0L) {
			this.mappingRefCount++;

			MappedView var10;
			try (MemoryStack stack = MemoryStack.stackPush()) {
				PointerBuffer pointer = stack.callocPointer(1);
				VulkanUtils.crashIfFailure(this.device, Vma.vmaMapMemory(this.device.vma(), this.vmaAllocation, pointer), "Failed to map buffer");
				ByteBuffer byteBuffer = MemoryUtil.memByteBuffer(pointer.get(0) + offset, (int)length);
				var10 = new MappedView(this.slice(offset, length), byteBuffer, new Runnable() {
					private boolean closed;

					{
						Objects.requireNonNull(VulkanGpuBuffer.this);
						this.closed = false;
					}

					public void run() {
						if (!this.closed) {
							this.closed = true;
							VulkanGpuBuffer.this.mappingRefCount--;
							Vma.vmaUnmapMemory(VulkanGpuBuffer.this.device.vma(), VulkanGpuBuffer.this.vmaAllocation);
						}
					}
				});
			}

			return var10;
		} else {
			throw new IllegalArgumentException("Offset or length must be positive integer values");
		}
	}
}