package com.mojang.blaze3d.opengl;

import com.mojang.blaze3d.buffers.GpuBufferSlice;
import com.mojang.blaze3d.vertex.VertexFormat;
import com.mojang.blaze3d.vertex.VertexFormatElement;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
import net.minecraft.util.VisibleForDebug;
import org.jspecify.annotations.Nullable;
import org.lwjgl.opengl.ARBVertexAttribBinding;
import org.lwjgl.opengl.GL33C;
import org.lwjgl.opengl.GLCapabilities;

public abstract class VertexArrayCache {
	public static VertexArrayCache create(final GLCapabilities capabilities, final GlDebugLabel debugLabels, final Set<String> enabledExtensions) {
		if (capabilities.GL_ARB_vertex_attrib_binding && GlDevice.USE_GL_ARB_vertex_attrib_binding) {
			enabledExtensions.add("GL_ARB_vertex_attrib_binding");
			return new VertexArrayCache.Separate(debugLabels);
		} else {
			return new VertexArrayCache.Emulated(debugLabels);
		}
	}

	public abstract VertexArrayCache.VertexArray bindVertexArray(
		final VertexFormat[] vertexBindings, final GpuBufferSlice[] vertexBuffers, @Nullable final VertexArrayCache.VertexArray lastBoundVertexArray
	);

	private static class Emulated extends VertexArrayCache {
		private final Map<List<VertexFormat>, VertexArrayCache.VertexArray> cache = new HashMap();
		private final GlDebugLabel debugLabels;

		public Emulated(final GlDebugLabel debugLabels) {
			this.debugLabels = debugLabels;
		}

		@Override
		public VertexArrayCache.VertexArray bindVertexArray(
			final VertexFormat[] vertexBindings, final GpuBufferSlice[] vertexBuffers, final VertexArrayCache.VertexArray lastBoundVertexArray
		) {
			List<VertexFormat> listBindings = Arrays.asList(vertexBindings);
			VertexArrayCache.VertexArray vertexArray = (VertexArrayCache.VertexArray)this.cache.get(listBindings);
			if (vertexArray == null) {
				int id = GlStateManager._glGenVertexArrays();
				GlStateManager._glBindVertexArray(id);
				setupCombinedAttributes(vertexBindings, true, vertexBuffers);
				VertexArrayCache.VertexArray vao = new VertexArrayCache.VertexArray(id, vertexBindings);
				this.debugLabels.applyLabel(vao);
				this.cache.put(listBindings, vao);
				return vao;
			} else {
				GlStateManager._glBindVertexArray(vertexArray.id);
				if (vertexArray != lastBoundVertexArray) {
					setupCombinedAttributes(vertexBindings, false, vertexBuffers);
				}

				return vertexArray;
			}
		}

		private static void setupCombinedAttributes(final VertexFormat[] vertexBindings, final boolean enable, final GpuBufferSlice[] vertexBuffers) {
			int attributeIndex = 0;

			for (int i = 0; i < vertexBindings.length; i++) {
				VertexFormat vertexBinding = vertexBindings[i];
				if (vertexBinding != null) {
					GlBuffer buffer = (GlBuffer)vertexBuffers[i].buffer();
					GlStateManager._glBindBuffer(34962, buffer.handle);
					int vertexSize = vertexBinding.getVertexSize();

					for (VertexFormatElement element : vertexBinding.getElements()) {
						long totalOffset = vertexBuffers[i].offset() + element.offset();
						int glExternalId = GlConst.toGlExternalId(element.format());
						int glType = GlConst.toGlType(element.format());
						boolean isIntegerFormat = GlConst.isGlFormatInteger(glExternalId);
						boolean isNormalizedFormat = GlConst.isFormatNormalized(element.format());
						int channelCount = GlConst.glFormatChannelCount(glExternalId);
						if (enable) {
							GlStateManager._enableVertexAttribArray(attributeIndex);
						}

						if (isIntegerFormat) {
							GlStateManager._vertexAttribIPointer(attributeIndex, channelCount, glType, vertexSize, totalOffset);
						} else {
							GlStateManager._vertexAttribPointer(attributeIndex, channelCount, glType, isNormalizedFormat, vertexSize, totalOffset);
						}

						GL33C.glVertexAttribDivisor(attributeIndex, vertexBinding.getStepRate());
						attributeIndex++;
					}
				}
			}
		}
	}

	private static class Separate extends VertexArrayCache {
		private final Map<List<VertexFormat>, VertexArrayCache.VertexArray> cache = new HashMap();
		private final GlDebugLabel debugLabels;
		private final boolean needsMesaWorkaround;

		public Separate(final GlDebugLabel debugLabels) {
			this.debugLabels = debugLabels;
			if ("Mesa".equals(GlStateManager._getString(7936))) {
				String version = GlStateManager._getString(7938);
				this.needsMesaWorkaround = version.contains("25.0.0") || version.contains("25.0.1") || version.contains("25.0.2");
			} else {
				this.needsMesaWorkaround = false;
			}
		}

		@Override
		public VertexArrayCache.VertexArray bindVertexArray(
			final VertexFormat[] vertexBindings, final GpuBufferSlice[] vertexBuffers, final VertexArrayCache.VertexArray lastBoundVertexArray
		) {
			List<VertexFormat> listBindings = Arrays.asList(vertexBindings);
			VertexArrayCache.VertexArray vertexArray = (VertexArrayCache.VertexArray)this.cache.get(listBindings);
			if (vertexArray == null) {
				int id = GlStateManager._glGenVertexArrays();
				GlStateManager._glBindVertexArray(id);
				int attribLocation = 0;

				for (int i = 0; i < vertexBindings.length; i++) {
					VertexFormat vertexBinding = vertexBindings[i];
					if (vertexBinding != null) {
						for (VertexFormatElement element : vertexBinding.getElements()) {
							if (element != null) {
								GlStateManager._enableVertexAttribArray(attribLocation);
								int glExternalId = GlConst.toGlExternalId(element.format());
								int glType = GlConst.toGlType(element.format());
								boolean isIntegerFormat = GlConst.isGlFormatInteger(glExternalId);
								boolean isNormalizedFormat = GlConst.isFormatNormalized(element.format());
								int channelCount = GlConst.glFormatChannelCount(glExternalId);
								if (isIntegerFormat) {
									ARBVertexAttribBinding.glVertexAttribIFormat(attribLocation, channelCount, glType, element.offset());
								} else {
									ARBVertexAttribBinding.glVertexAttribFormat(attribLocation, channelCount, glType, isNormalizedFormat, element.offset());
								}

								ARBVertexAttribBinding.glVertexAttribBinding(attribLocation, i);
								attribLocation++;
							}
						}

						ARBVertexAttribBinding.glVertexBindingDivisor(i, vertexBinding.getStepRate());
					}
				}

				for (int ix = 0; ix < vertexBuffers.length; ix++) {
					GpuBufferSlice vertexBufferSlice = vertexBuffers[ix];
					if (vertexBufferSlice != null) {
						GlBuffer vertexBuffer = (GlBuffer)vertexBufferSlice.buffer();
						ARBVertexAttribBinding.glBindVertexBuffer(ix, vertexBuffer.handle, vertexBufferSlice.offset(), vertexBindings[ix].getVertexSize());
					}
				}

				VertexArrayCache.VertexArray vao = new VertexArrayCache.VertexArray(id, vertexBindings);
				this.debugLabels.applyLabel(vao);
				this.cache.put(listBindings, vao);
				return vao;
			} else {
				GlStateManager._glBindVertexArray(vertexArray.id);
				if (vertexArray != lastBoundVertexArray) {
					for (int ixx = 0; ixx < vertexBuffers.length; ixx++) {
						GpuBufferSlice vertexBufferSlice = vertexBuffers[ixx];
						if (vertexBufferSlice != null) {
							GlBuffer vertexBuffer = (GlBuffer)vertexBufferSlice.buffer();
							if (this.needsMesaWorkaround) {
								ARBVertexAttribBinding.glBindVertexBuffer(ixx, 0, 0L, 0);
							}

							ARBVertexAttribBinding.glBindVertexBuffer(ixx, vertexBuffer.handle, vertexBufferSlice.offset(), vertexBindings[ixx].getVertexSize());
						}
					}
				}

				return vertexArray;
			}
		}
	}

	public static class VertexArray {
		@VisibleForDebug
		final int id;
		@VisibleForDebug
		final String formatName;

		private VertexArray(final int id, final VertexFormat[] vertexBindings) {
			this.id = id;
			this.formatName = (String)Arrays.stream(vertexBindings).filter(Objects::nonNull).map(VertexFormat::toString).collect(Collectors.joining(", "));
		}
	}
}