package net.minecraft.client.renderer;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableMap.Builder;
import com.google.gson.JsonElement;
import com.google.gson.JsonParseException;
import com.google.gson.JsonParser;
import com.google.gson.JsonSyntaxException;
import com.mojang.blaze3d.preprocessor.GlslPreprocessor;
import com.mojang.blaze3d.shaders.CompiledShader;
import com.mojang.blaze3d.shaders.CompiledShader.Type;
import com.mojang.blaze3d.systems.RenderSystem;
import com.mojang.logging.LogUtils;
import com.mojang.serialization.JsonOps;
import it.unimi.dsi.fastutil.objects.ObjectArraySet;
import java.io.IOException;
import java.io.Reader;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.Map.Entry;
import java.util.function.Consumer;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.FileUtil;
import net.minecraft.ResourceLocationException;
import net.minecraft.client.renderer.texture.TextureManager;
import net.minecraft.resources.FileToIdConverter;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.packs.resources.Resource;
import net.minecraft.server.packs.resources.ResourceManager;
import net.minecraft.server.packs.resources.ResourceProvider;
import net.minecraft.server.packs.resources.SimplePreparableReloadListener;
import net.minecraft.util.profiling.ProfilerFiller;
import org.apache.commons.io.IOUtils;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;

@Environment(EnvType.CLIENT)
public class ShaderManager extends SimplePreparableReloadListener<ShaderManager.Configs> implements AutoCloseable {
	static final Logger LOGGER = LogUtils.getLogger();
	public static final String SHADER_PATH = "shaders";
	public static final String SHADER_INCLUDE_PATH = "shaders/include/";
	private static final FileToIdConverter PROGRAM_ID_CONVERTER = FileToIdConverter.json("shaders");
	private static final FileToIdConverter POST_CHAIN_ID_CONVERTER = FileToIdConverter.json("post_effect");
	public static final int MAX_LOG_LENGTH = 32768;
	final TextureManager textureManager;
	private final Consumer<Exception> recoveryHandler;
	private ShaderManager.CompilationCache compilationCache = new ShaderManager.CompilationCache(ShaderManager.Configs.EMPTY);

	public ShaderManager(TextureManager textureManager, Consumer<Exception> recoveryHandler) {
		this.textureManager = textureManager;
		this.recoveryHandler = recoveryHandler;
	}

	protected ShaderManager.Configs prepare(ResourceManager resourceManager, ProfilerFiller profilerFiller) {
		Builder<ResourceLocation, ShaderProgramConfig> builder = ImmutableMap.builder();
		Builder<ShaderManager.ShaderSourceKey, String> builder2 = ImmutableMap.builder();
		Map<ResourceLocation, Resource> map = resourceManager.listResources("shaders", resourceLocation -> isProgram(resourceLocation) || isShader(resourceLocation));

		for (Entry<ResourceLocation, Resource> entry : map.entrySet()) {
			ResourceLocation resourceLocation = (ResourceLocation)entry.getKey();
			Type type = Type.byLocation(resourceLocation);
			if (type != null) {
				loadShader(resourceLocation, (Resource)entry.getValue(), type, map, builder2);
			} else if (isProgram(resourceLocation)) {
				loadProgram(resourceLocation, (Resource)entry.getValue(), builder);
			}
		}

		Builder<ResourceLocation, PostChainConfig> builder3 = ImmutableMap.builder();

		for (Entry<ResourceLocation, Resource> entry2 : POST_CHAIN_ID_CONVERTER.listMatchingResources(resourceManager).entrySet()) {
			loadPostChain((ResourceLocation)entry2.getKey(), (Resource)entry2.getValue(), builder3);
		}

		return new ShaderManager.Configs(builder.build(), builder2.build(), builder3.build());
	}

	private static void loadShader(
		ResourceLocation location, Resource shader, Type type, Map<ResourceLocation, Resource> shaderResources, Builder<ShaderManager.ShaderSourceKey, String> output
	) {
		ResourceLocation resourceLocation = type.idConverter().fileToId(location);
		GlslPreprocessor glslPreprocessor = createPreprocessor(shaderResources, location);

		try {
			Reader reader = shader.openAsReader();

			try {
				String string = IOUtils.toString(reader);
				output.put(new ShaderManager.ShaderSourceKey(resourceLocation, type), String.join("", glslPreprocessor.process(string)));
			} catch (Throwable var11) {
				if (reader != null) {
					try {
						reader.close();
					} catch (Throwable var10) {
						var11.addSuppressed(var10);
					}
				}

				throw var11;
			}

			if (reader != null) {
				reader.close();
			}
		} catch (IOException var12) {
			LOGGER.error("Failed to load shader source at {}", location, var12);
		}
	}

	private static GlslPreprocessor createPreprocessor(Map<ResourceLocation, Resource> shaderResources, ResourceLocation shaderLocation) {
		final ResourceLocation resourceLocation = shaderLocation.withPath(FileUtil::getFullResourcePath);
		return new GlslPreprocessor() {
			private final Set<ResourceLocation> importedLocations = new ObjectArraySet<>();

			@Override
			public String applyImport(boolean useFullPath, String directory) {
				ResourceLocation resourceLocationx;
				try {
					if (useFullPath) {
						resourceLocationx = resourceLocation.withPath((UnaryOperator<String>)(string2 -> FileUtil.normalizeResourcePath(string2 + directory)));
					} else {
						resourceLocationx = ResourceLocation.parse(directory).withPrefix("shaders/include/");
					}
				} catch (ResourceLocationException var8) {
					ShaderManager.LOGGER.error("Malformed GLSL import {}: {}", directory, var8.getMessage());
					return "#error " + var8.getMessage();
				}

				if (!this.importedLocations.add(resourceLocationx)) {
					return null;
				} else {
					try {
						Reader reader = ((Resource)shaderResources.get(resourceLocationx)).openAsReader();

						String var5;
						try {
							var5 = IOUtils.toString(reader);
						} catch (Throwable var9) {
							if (reader != null) {
								try {
									reader.close();
								} catch (Throwable var7) {
									var9.addSuppressed(var7);
								}
							}

							throw var9;
						}

						if (reader != null) {
							reader.close();
						}

						return var5;
					} catch (IOException var10) {
						ShaderManager.LOGGER.error("Could not open GLSL import {}: {}", resourceLocationx, var10.getMessage());
						return "#error " + var10.getMessage();
					}
				}
			}
		};
	}

	private static void loadProgram(ResourceLocation location, Resource resource, Builder<ResourceLocation, ShaderProgramConfig> output) {
		ResourceLocation resourceLocation = PROGRAM_ID_CONVERTER.fileToId(location);

		try {
			Reader reader = resource.openAsReader();

			try {
				JsonElement jsonElement = JsonParser.parseReader(reader);
				ShaderProgramConfig shaderProgramConfig = ShaderProgramConfig.CODEC.parse(JsonOps.INSTANCE, jsonElement).getOrThrow(JsonSyntaxException::new);
				output.put(resourceLocation, shaderProgramConfig);
			} catch (Throwable var8) {
				if (reader != null) {
					try {
						reader.close();
					} catch (Throwable var7) {
						var8.addSuppressed(var7);
					}
				}

				throw var8;
			}

			if (reader != null) {
				reader.close();
			}
		} catch (JsonParseException | IOException var9) {
			LOGGER.error("Failed to parse shader config at {}", location, var9);
		}
	}

	private static void loadPostChain(ResourceLocation location, Resource postChain, Builder<ResourceLocation, PostChainConfig> output) {
		ResourceLocation resourceLocation = POST_CHAIN_ID_CONVERTER.fileToId(location);

		try {
			Reader reader = postChain.openAsReader();

			try {
				JsonElement jsonElement = JsonParser.parseReader(reader);
				output.put(resourceLocation, PostChainConfig.CODEC.parse(JsonOps.INSTANCE, jsonElement).getOrThrow(JsonSyntaxException::new));
			} catch (Throwable var8) {
				if (reader != null) {
					try {
						reader.close();
					} catch (Throwable var7) {
						var8.addSuppressed(var7);
					}
				}

				throw var8;
			}

			if (reader != null) {
				reader.close();
			}
		} catch (JsonParseException | IOException var9) {
			LOGGER.error("Failed to parse post chain at {}", location, var9);
		}
	}

	private static boolean isProgram(ResourceLocation location) {
		return location.getPath().endsWith(".json");
	}

	private static boolean isShader(ResourceLocation location) {
		return Type.byLocation(location) != null || location.getPath().endsWith(".glsl");
	}

	protected void apply(ShaderManager.Configs configs, ResourceManager resourceManager, ProfilerFiller profilerFiller) {
		ShaderManager.CompilationCache compilationCache = new ShaderManager.CompilationCache(configs);
		Map<ShaderProgram, ShaderManager.CompilationException> map = new HashMap();
		Set<ShaderProgram> set = new HashSet(CoreShaders.getProgramsToPreload());

		for (PostChainConfig postChainConfig : configs.postChains.values()) {
			for (PostChainConfig.Pass pass : postChainConfig.passes()) {
				set.add(pass.program());
			}
		}

		for (ShaderProgram shaderProgram : set) {
			try {
				compilationCache.programs.put(shaderProgram, Optional.of(compilationCache.compileProgram(shaderProgram)));
			} catch (ShaderManager.CompilationException var11) {
				map.put(shaderProgram, var11);
			}
		}

		if (!map.isEmpty()) {
			compilationCache.close();
			throw new RuntimeException(
				"Failed to load required shader programs:\n"
					+ (String)map.entrySet()
						.stream()
						.map(entry -> " - " + entry.getKey() + ": " + ((ShaderManager.CompilationException)entry.getValue()).getMessage())
						.collect(Collectors.joining("\n"))
			);
		} else {
			this.compilationCache.close();
			this.compilationCache = compilationCache;
		}
	}

	@Override
	public String getName() {
		return "Shader Loader";
	}

	private void tryTriggerRecovery(Exception exception) {
		if (!this.compilationCache.triggeredRecovery) {
			this.recoveryHandler.accept(exception);
			this.compilationCache.triggeredRecovery = true;
		}
	}

	public void preloadForStartup(ResourceProvider resourceProvider, ShaderProgram... programs) throws IOException, ShaderManager.CompilationException {
		for (ShaderProgram shaderProgram : programs) {
			Resource resource = resourceProvider.getResourceOrThrow(PROGRAM_ID_CONVERTER.idToFile(shaderProgram.configId()));
			Reader reader = resource.openAsReader();

			try {
				JsonElement jsonElement = JsonParser.parseReader(reader);
				ShaderProgramConfig shaderProgramConfig = ShaderProgramConfig.CODEC.parse(JsonOps.INSTANCE, jsonElement).getOrThrow(JsonSyntaxException::new);
				ShaderDefines shaderDefines = shaderProgramConfig.defines().withOverrides(shaderProgram.defines());
				CompiledShader compiledShader = this.preloadShader(resourceProvider, shaderProgramConfig.vertex(), Type.VERTEX, shaderDefines);
				CompiledShader compiledShader2 = this.preloadShader(resourceProvider, shaderProgramConfig.fragment(), Type.FRAGMENT, shaderDefines);
				CompiledShaderProgram compiledShaderProgram = linkProgram(shaderProgram, shaderProgramConfig, compiledShader, compiledShader2);
				this.compilationCache.programs.put(shaderProgram, Optional.of(compiledShaderProgram));
			} catch (Throwable var16) {
				if (reader != null) {
					try {
						reader.close();
					} catch (Throwable var15) {
						var16.addSuppressed(var15);
					}
				}

				throw var16;
			}

			if (reader != null) {
				reader.close();
			}
		}
	}

	private CompiledShader preloadShader(ResourceProvider resourceProvider, ResourceLocation shader, Type type, ShaderDefines defines) throws IOException, ShaderManager.CompilationException {
		ResourceLocation resourceLocation = type.idConverter().idToFile(shader);
		Reader reader = resourceProvider.getResourceOrThrow(resourceLocation).openAsReader();

		CompiledShader var10;
		try {
			String string = IOUtils.toString(reader);
			String string2 = GlslPreprocessor.injectDefines(string, defines);
			CompiledShader compiledShader = CompiledShader.compile(shader, type, string2);
			this.compilationCache.shaders.put(new ShaderManager.ShaderCompilationKey(shader, type, defines), compiledShader);
			var10 = compiledShader;
		} catch (Throwable var12) {
			if (reader != null) {
				try {
					reader.close();
				} catch (Throwable var11) {
					var12.addSuppressed(var11);
				}
			}

			throw var12;
		}

		if (reader != null) {
			reader.close();
		}

		return var10;
	}

	@Nullable
	public CompiledShaderProgram getProgram(ShaderProgram program) {
		try {
			return this.compilationCache.getOrCompileProgram(program);
		} catch (ShaderManager.CompilationException var3) {
			LOGGER.error("Failed to load shader program: {}", program, var3);
			this.compilationCache.programs.put(program, Optional.empty());
			this.tryTriggerRecovery(var3);
			return null;
		}
	}

	public CompiledShaderProgram getProgramForLoading(ShaderProgram program) throws ShaderManager.CompilationException {
		CompiledShaderProgram compiledShaderProgram = this.compilationCache.getOrCompileProgram(program);
		if (compiledShaderProgram == null) {
			throw new ShaderManager.CompilationException("Shader '" + program + "' could not be found");
		} else {
			return compiledShaderProgram;
		}
	}

	static CompiledShaderProgram linkProgram(ShaderProgram program, ShaderProgramConfig config, CompiledShader vertexShader, CompiledShader fragmentShader) throws ShaderManager.CompilationException {
		CompiledShaderProgram compiledShaderProgram = CompiledShaderProgram.link(vertexShader, fragmentShader, program.vertexFormat());
		compiledShaderProgram.setupUniforms(config.uniforms(), config.samplers());
		return compiledShaderProgram;
	}

	@Nullable
	public PostChain getPostChain(ResourceLocation id, Set<ResourceLocation> externalTargets) {
		try {
			return this.compilationCache.getOrLoadPostChain(id, externalTargets);
		} catch (ShaderManager.CompilationException var4) {
			LOGGER.error("Failed to load post chain: {}", id, var4);
			this.compilationCache.postChains.put(id, Optional.empty());
			this.tryTriggerRecovery(var4);
			return null;
		}
	}

	public void close() {
		this.compilationCache.close();
	}

	@Environment(EnvType.CLIENT)
	class CompilationCache implements AutoCloseable {
		private final ShaderManager.Configs configs;
		final Map<ShaderProgram, Optional<CompiledShaderProgram>> programs = new HashMap();
		final Map<ShaderManager.ShaderCompilationKey, CompiledShader> shaders = new HashMap();
		final Map<ResourceLocation, Optional<PostChain>> postChains = new HashMap();
		boolean triggeredRecovery;

		CompilationCache(final ShaderManager.Configs configs) {
			this.configs = configs;
		}

		@Nullable
		public CompiledShaderProgram getOrCompileProgram(ShaderProgram program) throws ShaderManager.CompilationException {
			Optional<CompiledShaderProgram> optional = (Optional<CompiledShaderProgram>)this.programs.get(program);
			if (optional != null) {
				return (CompiledShaderProgram)optional.orElse(null);
			} else {
				CompiledShaderProgram compiledShaderProgram = this.compileProgram(program);
				this.programs.put(program, Optional.of(compiledShaderProgram));
				return compiledShaderProgram;
			}
		}

		CompiledShaderProgram compileProgram(ShaderProgram program) throws ShaderManager.CompilationException {
			ShaderProgramConfig shaderProgramConfig = (ShaderProgramConfig)this.configs.programs.get(program.configId());
			if (shaderProgramConfig == null) {
				throw new ShaderManager.CompilationException("Could not find program with id: " + program.configId());
			} else {
				ShaderDefines shaderDefines = shaderProgramConfig.defines().withOverrides(program.defines());
				CompiledShader compiledShader = this.getOrCompileShader(shaderProgramConfig.vertex(), Type.VERTEX, shaderDefines);
				CompiledShader compiledShader2 = this.getOrCompileShader(shaderProgramConfig.fragment(), Type.FRAGMENT, shaderDefines);
				return ShaderManager.linkProgram(program, shaderProgramConfig, compiledShader, compiledShader2);
			}
		}

		private CompiledShader getOrCompileShader(ResourceLocation id, Type type, ShaderDefines defines) throws ShaderManager.CompilationException {
			ShaderManager.ShaderCompilationKey shaderCompilationKey = new ShaderManager.ShaderCompilationKey(id, type, defines);
			CompiledShader compiledShader = (CompiledShader)this.shaders.get(shaderCompilationKey);
			if (compiledShader == null) {
				compiledShader = this.compileShader(shaderCompilationKey);
				this.shaders.put(shaderCompilationKey, compiledShader);
			}

			return compiledShader;
		}

		private CompiledShader compileShader(ShaderManager.ShaderCompilationKey compilationKey) throws ShaderManager.CompilationException {
			String string = (String)this.configs.shaderSources.get(new ShaderManager.ShaderSourceKey(compilationKey.id, compilationKey.type));
			if (string == null) {
				throw new ShaderManager.CompilationException("Could not find shader: " + compilationKey);
			} else {
				String string2 = GlslPreprocessor.injectDefines(string, compilationKey.defines);
				return CompiledShader.compile(compilationKey.id, compilationKey.type, string2);
			}
		}

		@Nullable
		public PostChain getOrLoadPostChain(ResourceLocation name, Set<ResourceLocation> externalTargets) throws ShaderManager.CompilationException {
			Optional<PostChain> optional = (Optional<PostChain>)this.postChains.get(name);
			if (optional != null) {
				return (PostChain)optional.orElse(null);
			} else {
				PostChain postChain = this.loadPostChain(name, externalTargets);
				this.postChains.put(name, Optional.of(postChain));
				return postChain;
			}
		}

		private PostChain loadPostChain(ResourceLocation name, Set<ResourceLocation> externalTargets) throws ShaderManager.CompilationException {
			PostChainConfig postChainConfig = (PostChainConfig)this.configs.postChains.get(name);
			if (postChainConfig == null) {
				throw new ShaderManager.CompilationException("Could not find post chain with id: " + name);
			} else {
				return PostChain.load(postChainConfig, ShaderManager.this.textureManager, ShaderManager.this, externalTargets);
			}
		}

		public void close() {
			RenderSystem.assertOnRenderThread();
			this.programs.values().forEach(optional -> optional.ifPresent(CompiledShaderProgram::close));
			this.shaders.values().forEach(CompiledShader::close);
			this.programs.clear();
			this.shaders.clear();
			this.postChains.clear();
		}
	}

	@Environment(EnvType.CLIENT)
	public static class CompilationException extends Exception {
		public CompilationException(String message) {
			super(message);
		}
	}

	@Environment(EnvType.CLIENT)
	public record Configs(
		Map<ResourceLocation, ShaderProgramConfig> programs,
		Map<ShaderManager.ShaderSourceKey, String> shaderSources,
		Map<ResourceLocation, PostChainConfig> postChains
	) {
		public static final ShaderManager.Configs EMPTY = new ShaderManager.Configs(Map.of(), Map.of(), Map.of());
	}

	@Environment(EnvType.CLIENT)
	record ShaderCompilationKey(ResourceLocation id, Type type, ShaderDefines defines) {

		public String toString() {
			String string = this.id + " (" + this.type + ")";
			return !this.defines.isEmpty() ? string + " with " + this.defines : string;
		}
	}

	@Environment(EnvType.CLIENT)
	record ShaderSourceKey(ResourceLocation id, Type type) {
		public String toString() {
			return this.id + " (" + this.type + ")";
		}
	}
}