package com.mojang.blaze3d.vertex;

import com.mojang.blaze3d.platform.GlStateManager;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Stream;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import org.jetbrains.annotations.Nullable;

@Environment(EnvType.CLIENT)
public record VertexFormatElement(int id, int index, VertexFormatElement.Type type, VertexFormatElement.Usage usage, int count) {
	public static final int MAX_COUNT = 32;
	private static final VertexFormatElement[] BY_ID = new VertexFormatElement[32];
	private static final List<VertexFormatElement> ELEMENTS = new ArrayList(32);
	public static final VertexFormatElement POSITION = register(0, 0, VertexFormatElement.Type.FLOAT, VertexFormatElement.Usage.POSITION, 3);
	public static final VertexFormatElement COLOR = register(1, 0, VertexFormatElement.Type.UBYTE, VertexFormatElement.Usage.COLOR, 4);
	public static final VertexFormatElement UV0 = register(2, 0, VertexFormatElement.Type.FLOAT, VertexFormatElement.Usage.UV, 2);
	public static final VertexFormatElement UV = UV0;
	public static final VertexFormatElement UV1 = register(3, 1, VertexFormatElement.Type.SHORT, VertexFormatElement.Usage.UV, 2);
	public static final VertexFormatElement UV2 = register(4, 2, VertexFormatElement.Type.SHORT, VertexFormatElement.Usage.UV, 2);
	public static final VertexFormatElement NORMAL = register(5, 0, VertexFormatElement.Type.BYTE, VertexFormatElement.Usage.NORMAL, 3);

	public VertexFormatElement(int id, int index, VertexFormatElement.Type type, VertexFormatElement.Usage usage, int count) {
		if (id < 0 || id >= BY_ID.length) {
			throw new IllegalArgumentException("Element ID must be in range [0; " + BY_ID.length + ")");
		} else if (!this.supportsUsage(index, usage)) {
			throw new IllegalStateException("Multiple vertex elements of the same type other than UVs are not supported");
		} else {
			this.id = id;
			this.index = index;
			this.type = type;
			this.usage = usage;
			this.count = count;
		}
	}

	public static VertexFormatElement register(int id, int index, VertexFormatElement.Type type, VertexFormatElement.Usage usage, int count) {
		VertexFormatElement vertexFormatElement = new VertexFormatElement(id, index, type, usage, count);
		if (BY_ID[id] != null) {
			throw new IllegalArgumentException("Duplicate element registration for: " + id);
		} else {
			BY_ID[id] = vertexFormatElement;
			ELEMENTS.add(vertexFormatElement);
			return vertexFormatElement;
		}
	}

	private boolean supportsUsage(int index, VertexFormatElement.Usage usage) {
		return index == 0 || usage == VertexFormatElement.Usage.UV;
	}

	public String toString() {
		return this.count + "," + this.usage + "," + this.type + " (" + this.id + ")";
	}

	public int mask() {
		return 1 << this.id;
	}

	public int byteSize() {
		return this.type.size() * this.count;
	}

	public void setupBufferState(int stateIndex, long offset, int stride) {
		this.usage.setupState.setupBufferState(this.count, this.type.glType(), stride, offset, stateIndex);
	}

	@Nullable
	public static VertexFormatElement byId(int id) {
		return BY_ID[id];
	}

	public static Stream<VertexFormatElement> elementsFromMask(int mask) {
		return ELEMENTS.stream().filter(vertexFormatElement -> vertexFormatElement != null && (mask & vertexFormatElement.mask()) != 0);
	}

	@Environment(EnvType.CLIENT)
	public static enum Type {
		FLOAT(4, "Float", 5126),
		UBYTE(1, "Unsigned Byte", 5121),
		BYTE(1, "Byte", 5120),
		USHORT(2, "Unsigned Short", 5123),
		SHORT(2, "Short", 5122),
		UINT(4, "Unsigned Int", 5125),
		INT(4, "Int", 5124);

		private final int size;
		private final String name;
		private final int glType;

		private Type(final int size, final String name, final int glType) {
			this.size = size;
			this.name = name;
			this.glType = glType;
		}

		public int size() {
			return this.size;
		}

		public int glType() {
			return this.glType;
		}

		public String toString() {
			return this.name;
		}
	}

	@Environment(EnvType.CLIENT)
	public static enum Usage {
		POSITION("Position", (i, j, k, l, m) -> GlStateManager._vertexAttribPointer(m, i, j, false, k, l)),
		NORMAL("Normal", (i, j, k, l, m) -> GlStateManager._vertexAttribPointer(m, i, j, true, k, l)),
		COLOR("Vertex Color", (i, j, k, l, m) -> GlStateManager._vertexAttribPointer(m, i, j, true, k, l)),
		UV("UV", (i, j, k, l, m) -> {
			if (j == 5126) {
				GlStateManager._vertexAttribPointer(m, i, j, false, k, l);
			} else {
				GlStateManager._vertexAttribIPointer(m, i, j, k, l);
			}
		}),
		GENERIC("Generic", (i, j, k, l, m) -> GlStateManager._vertexAttribPointer(m, i, j, false, k, l));

		private final String name;
		final VertexFormatElement.Usage.SetupState setupState;

		private Usage(final String name, final VertexFormatElement.Usage.SetupState setupState) {
			this.name = name;
			this.setupState = setupState;
		}

		public String toString() {
			return this.name;
		}

		@FunctionalInterface
		@Environment(EnvType.CLIENT)
		interface SetupState {
			void setupBufferState(int i, int j, int k, long l, int m);
		}
	}
}