package com.mojang.blaze3d.vulkan.glsl;

import com.mojang.blaze3d.vulkan.VulkanBindGroupLayout;
import com.mojang.blaze3d.vulkan.VulkanDevice;
import com.mojang.blaze3d.vulkan.VulkanUtils;
import java.nio.ByteBuffer;
import java.nio.IntBuffer;
import java.nio.LongBuffer;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import org.jspecify.annotations.Nullable;
import org.lwjgl.PointerBuffer;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.system.MemoryUtil;
import org.lwjgl.util.spvc.Spvc;
import org.lwjgl.util.spvc.SpvcReflectedResource;
import org.lwjgl.util.spvc.SpvcReflectedResource.Buffer;
import org.lwjgl.vulkan.VK12;
import org.lwjgl.vulkan.VkShaderModuleCreateInfo;

@Environment(EnvType.CLIENT)
public record IntermediaryShaderModule(
	String name, ByteBuffer spirv, List<SpvUniformBuffer> uniformBuffers, List<SpvSampler> samplers, List<SpvVariable> outputs, List<SpvVariable> inputs
) implements AutoCloseable {
	public static final IntermediaryShaderModule INVALID = new IntermediaryShaderModule(
		"invalid", ByteBuffer.allocate(1), new ArrayList(), new ArrayList(), new ArrayList(), new ArrayList()
	);

	public static IntermediaryShaderModule createFromSpirv(final String filename, final ByteBuffer spirv) throws ShaderCompileException {
		List<SpvUniformBuffer> uniformBuffers = new ArrayList();
		List<SpvSampler> samplers = new ArrayList();
		List<SpvVariable> outputs = new ArrayList();
		List<SpvVariable> inputs = new ArrayList();

		try (MemoryStack stack = MemoryStack.stackPush()) {
			PointerBuffer pointer = stack.callocPointer(1);
			IntBuffer intReturnBuffer = stack.callocInt(1);
			throwIfError(Spvc.spvc_context_create(pointer), "Couldn't create spvc context");
			long context = pointer.get(0);

			try {
				throwIfError(Spvc.spvc_context_parse_spirv(context, spirv.asIntBuffer(), spirv.remaining() / 4, pointer), "Couldn't parse spirv");
				long ir = pointer.get(0);
				throwIfError(Spvc.spvc_context_create_compiler(context, 0, ir, 1, pointer), "Couldn't create compiler");
				long compiler = pointer.get(0);
				throwIfError(Spvc.spvc_compiler_create_shader_resources(compiler, pointer), "Couldn't create resource list");
				long spvcResources = pointer.get(0);
				PointerBuffer countPointer = stack.callocPointer(1);
				throwIfError(Spvc.spvc_resources_get_resource_list_for_type(spvcResources, 1, pointer, countPointer), "Couldn't list uniform buffers");
				long spvcList = pointer.get(0);
				long spvcCount = countPointer.get(0);
				Buffer resources = SpvcReflectedResource.create(spvcList, (int)spvcCount);

				for (int i = 0; i < spvcCount; i++) {
					SpvcReflectedResource resource = resources.get(i);
					String name = resource.nameString();
					int bindingOffset = getDecorationOffset(compiler, resource, 33, intReturnBuffer);
					uniformBuffers.add(new SpvUniformBuffer(name, bindingOffset));
				}

				throwIfError(Spvc.spvc_resources_get_resource_list_for_type(spvcResources, 7, pointer, countPointer), "Couldn't list sampled images");
				spvcList = pointer.get(0);
				spvcCount = countPointer.get(0);
				resources = SpvcReflectedResource.create(spvcList, (int)spvcCount);

				for (int i = 0; i < spvcCount; i++) {
					SpvcReflectedResource resource = resources.get(i);
					String name = resource.nameString();
					int bindingOffset = getDecorationOffset(compiler, resource, 33, intReturnBuffer);
					long typeHandle = Spvc.spvc_compiler_get_type_handle(compiler, resource.type_id());
					int dimension = Spvc.spvc_type_get_image_dimension(typeHandle);
					samplers.add(new SpvSampler(name, bindingOffset, dimension));
				}

				throwIfError(Spvc.spvc_resources_get_resource_list_for_type(spvcResources, 4, pointer, countPointer), "Couldn't list output variables");
				spvcList = pointer.get(0);
				spvcCount = countPointer.get(0);
				resources = SpvcReflectedResource.create(spvcList, (int)spvcCount);

				for (int i = 0; i < spvcCount; i++) {
					SpvcReflectedResource resource = resources.get(i);
					String name = resource.nameString();
					int bindingOffset = getDecorationOffset(compiler, resource, 30, intReturnBuffer);
					outputs.add(new SpvVariable(name, bindingOffset));
				}

				throwIfError(Spvc.spvc_resources_get_resource_list_for_type(spvcResources, 3, pointer, countPointer), "Couldn't list input variables");
				spvcList = pointer.get(0);
				spvcCount = countPointer.get(0);
				resources = SpvcReflectedResource.create(spvcList, (int)spvcCount);

				for (int i = 0; i < spvcCount; i++) {
					SpvcReflectedResource resource = resources.get(i);
					String name = resource.nameString();
					int bindingOffset = getDecorationOffset(compiler, resource, 30, intReturnBuffer);
					inputs.add(new SpvVariable(name, bindingOffset));
				}
			} finally {
				Spvc.spvc_context_destroy(context);
			}
		}

		IntBuffer spvAsIntBuffer = spirv.asIntBuffer();

		for (int i = 0; i < outputs.size(); i++) {
			spvAsIntBuffer.put(((SpvVariable)outputs.get(i)).locationOffset(), i);
		}

		return new IntermediaryShaderModule(filename, spirv, uniformBuffers, samplers, outputs, inputs);
	}

	public void close() {
		MemoryUtil.memFree(this.spirv);
	}

	public void rebind(final List<String> inputVariables, final List<VulkanBindGroupLayout.Entry> entries) throws ShaderCompileException {
		IntBuffer spvAsIntBuffer = this.spirv.asIntBuffer();
		Set<String> remainingInputs = new HashSet();
		Set<String> remainingSamplers = new HashSet();
		Set<String> remainingUniformBuffers = new HashSet();

		for (SpvVariable input : this.inputs) {
			remainingInputs.add(input.name());
		}

		for (SpvUniformBuffer uniformBuffer : this.uniformBuffers) {
			remainingUniformBuffers.add(uniformBuffer.name());
		}

		for (SpvSampler sampler : this.samplers) {
			remainingSamplers.add(sampler.name());
		}

		String previousName = null;
		int attribLocation = 0;

		for (int i = 0; i < inputVariables.size(); i++) {
			String variableName = (String)inputVariables.get(i);
			SpvVariable inputVariable = this.getInputVariable(variableName);
			if (inputVariable != null) {
				if (!variableName.equals(previousName)) {
					spvAsIntBuffer.put(inputVariable.locationOffset(), attribLocation);
					remainingInputs.remove(variableName);
				}

				attribLocation++;
				previousName = variableName;
			}
		}

		for (int ix = 0; ix < entries.size(); ix++) {
			VulkanBindGroupLayout.Entry entry = (VulkanBindGroupLayout.Entry)entries.get(ix);
			switch (entry.type()) {
				case UNIFORM_BUFFER:
					SpvUniformBuffer ubo = this.getUniformBuffer(entry.name());
					if (ubo != null) {
						spvAsIntBuffer.put(ubo.bindingOffset(), ix);
						remainingUniformBuffers.remove(entry.name());
					}
					break;
				case SAMPLED_IMAGE:
					SpvSampler samplerx = this.getSampler(entry.name());
					if (samplerx != null) {
						if (samplerx.dimensions() != 1 && samplerx.dimensions() != 3) {
							throw new ShaderCompileException(
								"Unsupported texture dimensions '" + SpvcUtil.imageDimensionToString(samplerx.dimensions()) + "' for sampler " + entry.name()
							);
						}

						spvAsIntBuffer.put(samplerx.bindingOffset(), ix);
						remainingSamplers.remove(entry.name());
					}
					break;
				case TEXEL_BUFFER:
					SpvSampler sampler = this.getSampler(entry.name());
					if (sampler != null) {
						if (sampler.dimensions() != 5) {
							throw new ShaderCompileException(
								"Unsupported texel buffer dimensions '" + SpvcUtil.imageDimensionToString(sampler.dimensions()) + "' for sampler " + entry.name()
							);
						}

						spvAsIntBuffer.put(sampler.bindingOffset(), ix);
						remainingSamplers.remove(entry.name());
					}
			}
		}

		if (!remainingInputs.isEmpty()) {
			throw new ShaderCompileException("Shader expects input variables which are not being provided: " + remainingInputs);
		} else if (!remainingUniformBuffers.isEmpty()) {
			throw new ShaderCompileException("Shader expects uniform buffers which are not being provided: " + remainingUniformBuffers);
		} else if (!remainingSamplers.isEmpty()) {
			throw new ShaderCompileException("Shader expects samplers which are not being provided: " + remainingSamplers);
		}
	}

	public long createVulkanShaderModule(final VulkanDevice device) {
		long var5;
		try (MemoryStack stack = MemoryStack.stackPush()) {
			VkShaderModuleCreateInfo info = VkShaderModuleCreateInfo.calloc(stack).sType$Default().pCode(this.spirv);
			LongBuffer pointer = stack.callocLong(1);
			VulkanUtils.crashIfFailure(VK12.vkCreateShaderModule(device.vkDevice(), info, null, pointer), "Can't compile " + this.name);
			device.instance().debug().setObjectName(device.vkDevice(), 15, pointer.get(0), () -> this.name);
			var5 = pointer.get(0);
		}

		return var5;
	}

	@Nullable
	private SpvUniformBuffer getUniformBuffer(final String name) {
		for (SpvUniformBuffer ubo : this.uniformBuffers) {
			if (ubo.name().equals(name)) {
				return ubo;
			}
		}

		return null;
	}

	@Nullable
	private SpvSampler getSampler(final String name) {
		for (SpvSampler sampler : this.samplers) {
			if (sampler.name().equals(name)) {
				return sampler;
			}
		}

		return null;
	}

	@Nullable
	private SpvVariable getInputVariable(final String name) {
		for (SpvVariable variable : this.inputs) {
			if (variable.name().equals(name)) {
				return variable;
			}
		}

		return null;
	}

	private static void throwIfError(final int result, final String message) throws ShaderCompileException {
		if (result != 0) {
			String name = switch (result) {
				case -4 -> "SPVC_ERROR_INVALID_ARGUMENT";
				case -3 -> "SPVC_ERROR_OUT_OF_MEMORY";
				case -2 -> "SPVC_ERROR_UNSUPPORTED_SPIRV";
				case -1 -> "SPVC_ERROR_INVALID_SPIRV";
				default -> Integer.toString(result);
			};
			throw new ShaderCompileException(message + " (" + name + ")");
		}
	}

	private static int getDecorationOffset(final long compiler, final SpvcReflectedResource resource, final int decoration, final IntBuffer returnBuffer) throws ShaderCompileException {
		if (!Spvc.spvc_compiler_get_binary_offset_for_decoration(compiler, resource.id(), decoration, returnBuffer)) {
			throw new ShaderCompileException("Couldn't find byte offset for location decoration of " + resource.nameString());
		} else {
			return returnBuffer.get(0);
		}
	}
}