package net.minecraft.client.renderer;

import com.mojang.blaze3d.shaders.UniformType;
import com.mojang.blaze3d.systems.RenderPass;
import com.mojang.datafixers.util.Either;
import com.mojang.serialization.Codec;
import com.mojang.serialization.DataResult;
import com.mojang.serialization.codecs.RecordCodecBuilder;
import it.unimi.dsi.fastutil.objects.ObjectArraySet;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Stream;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.util.ExtraCodecs;

@Environment(EnvType.CLIENT)
public record PostChainConfig(Map<ResourceLocation, PostChainConfig.InternalTarget> internalTargets, List<PostChainConfig.Pass> passes) {
	public static final Codec<PostChainConfig> CODEC = RecordCodecBuilder.create(
		instance -> instance.group(
				Codec.unboundedMap(ResourceLocation.CODEC, PostChainConfig.InternalTarget.CODEC)
					.optionalFieldOf("targets", Map.of())
					.forGetter(PostChainConfig::internalTargets),
				PostChainConfig.Pass.CODEC.listOf().optionalFieldOf("passes", List.of()).forGetter(PostChainConfig::passes)
			)
			.apply(instance, PostChainConfig::new)
	);

	@Environment(EnvType.CLIENT)
	public record FixedSizedTarget(int width, int height) implements PostChainConfig.InternalTarget {
		public static final Codec<PostChainConfig.FixedSizedTarget> CODEC = RecordCodecBuilder.create(
			instance -> instance.group(
					ExtraCodecs.POSITIVE_INT.fieldOf("width").forGetter(PostChainConfig.FixedSizedTarget::width),
					ExtraCodecs.POSITIVE_INT.fieldOf("height").forGetter(PostChainConfig.FixedSizedTarget::height)
				)
				.apply(instance, PostChainConfig.FixedSizedTarget::new)
		);
	}

	@Environment(EnvType.CLIENT)
	public record FullScreenTarget() implements PostChainConfig.InternalTarget {
		public static final Codec<PostChainConfig.FullScreenTarget> CODEC = Codec.unit(PostChainConfig.FullScreenTarget::new);
	}

	@Environment(EnvType.CLIENT)
	public sealed interface Input permits PostChainConfig.TextureInput, PostChainConfig.TargetInput {
		Codec<PostChainConfig.Input> CODEC = Codec.xor(PostChainConfig.TextureInput.CODEC, PostChainConfig.TargetInput.CODEC)
			.xmap(either -> either.map(Function.identity(), Function.identity()), input -> {
				return switch (input) {
					case PostChainConfig.TextureInput textureInput -> Either.left(textureInput);
					case PostChainConfig.TargetInput targetInput -> Either.right(targetInput);
					default -> throw new MatchException(null, null);
				};
			});

		String samplerName();

		Set<ResourceLocation> referencedTargets();
	}

	@Environment(EnvType.CLIENT)
	public sealed interface InternalTarget permits PostChainConfig.FullScreenTarget, PostChainConfig.FixedSizedTarget {
		Codec<PostChainConfig.InternalTarget> CODEC = Codec.either(PostChainConfig.FixedSizedTarget.CODEC, PostChainConfig.FullScreenTarget.CODEC)
			.xmap(either -> either.map(Function.identity(), Function.identity()), internalTarget -> {
				return switch (internalTarget) {
					case PostChainConfig.FixedSizedTarget fixedSizedTarget -> Either.left(fixedSizedTarget);
					case PostChainConfig.FullScreenTarget fullScreenTarget -> Either.right(fullScreenTarget);
					default -> throw new MatchException(null, null);
				};
			});
	}

	@Environment(EnvType.CLIENT)
	public record Pass(
		ResourceLocation vertexShaderId,
		ResourceLocation fragmentShaderId,
		List<PostChainConfig.Input> inputs,
		ResourceLocation outputTarget,
		List<PostChainConfig.Uniform> uniforms
	) {
		private static final Codec<List<PostChainConfig.Input>> INPUTS_CODEC = PostChainConfig.Input.CODEC.listOf().validate(list -> {
			Set<String> set = new ObjectArraySet<>(list.size());

			for (PostChainConfig.Input input : list) {
				if (!set.add(input.samplerName())) {
					return DataResult.error(() -> "Encountered repeated sampler name: " + input.samplerName());
				}
			}

			return DataResult.success(list);
		});
		public static final Codec<PostChainConfig.Pass> CODEC = RecordCodecBuilder.create(
			instance -> instance.group(
					ResourceLocation.CODEC.fieldOf("vertex_shader").forGetter(PostChainConfig.Pass::vertexShaderId),
					ResourceLocation.CODEC.fieldOf("fragment_shader").forGetter(PostChainConfig.Pass::fragmentShaderId),
					INPUTS_CODEC.optionalFieldOf("inputs", List.of()).forGetter(PostChainConfig.Pass::inputs),
					ResourceLocation.CODEC.fieldOf("output").forGetter(PostChainConfig.Pass::outputTarget),
					PostChainConfig.Uniform.CODEC.listOf().optionalFieldOf("uniforms", List.of()).forGetter(PostChainConfig.Pass::uniforms)
				)
				.apply(instance, PostChainConfig.Pass::new)
		);

		public Stream<ResourceLocation> referencedTargets() {
			Stream<ResourceLocation> stream = this.inputs.stream().flatMap(input -> input.referencedTargets().stream());
			return Stream.concat(stream, Stream.of(this.outputTarget));
		}
	}

	@Environment(EnvType.CLIENT)
	public record TargetInput(String samplerName, ResourceLocation targetId, boolean useDepthBuffer, boolean bilinear) implements PostChainConfig.Input {
		public static final Codec<PostChainConfig.TargetInput> CODEC = RecordCodecBuilder.create(
			instance -> instance.group(
					Codec.STRING.fieldOf("sampler_name").forGetter(PostChainConfig.TargetInput::samplerName),
					ResourceLocation.CODEC.fieldOf("target").forGetter(PostChainConfig.TargetInput::targetId),
					Codec.BOOL.optionalFieldOf("use_depth_buffer", false).forGetter(PostChainConfig.TargetInput::useDepthBuffer),
					Codec.BOOL.optionalFieldOf("bilinear", false).forGetter(PostChainConfig.TargetInput::bilinear)
				)
				.apply(instance, PostChainConfig.TargetInput::new)
		);

		@Override
		public Set<ResourceLocation> referencedTargets() {
			return Set.of(this.targetId);
		}
	}

	@Environment(EnvType.CLIENT)
	public record TextureInput(String samplerName, ResourceLocation location, int width, int height, boolean bilinear) implements PostChainConfig.Input {
		public static final Codec<PostChainConfig.TextureInput> CODEC = RecordCodecBuilder.create(
			instance -> instance.group(
					Codec.STRING.fieldOf("sampler_name").forGetter(PostChainConfig.TextureInput::samplerName),
					ResourceLocation.CODEC.fieldOf("location").forGetter(PostChainConfig.TextureInput::location),
					ExtraCodecs.POSITIVE_INT.fieldOf("width").forGetter(PostChainConfig.TextureInput::width),
					ExtraCodecs.POSITIVE_INT.fieldOf("height").forGetter(PostChainConfig.TextureInput::height),
					Codec.BOOL.optionalFieldOf("bilinear", false).forGetter(PostChainConfig.TextureInput::bilinear)
				)
				.apply(instance, PostChainConfig.TextureInput::new)
		);

		@Override
		public Set<ResourceLocation> referencedTargets() {
			return Set.of();
		}
	}

	@Environment(EnvType.CLIENT)
	public record Uniform(String name, String type, Optional<List<Float>> values) {
		public static final Codec<PostChainConfig.Uniform> CODEC = RecordCodecBuilder.create(
			instance -> instance.group(
					Codec.STRING.fieldOf("name").forGetter(PostChainConfig.Uniform::name),
					Codec.STRING.fieldOf("type").forGetter(PostChainConfig.Uniform::type),
					Codec.FLOAT.sizeLimitedListOf(4).optionalFieldOf("values").forGetter(PostChainConfig.Uniform::values)
				)
				.apply(instance, PostChainConfig.Uniform::new)
		);

		public void setOnRenderPass(RenderPass renderPass) {
			UniformType uniformType = (UniformType)UniformType.CODEC.byName(this.type);
			if (!this.values.isEmpty() && uniformType != null && !((List)this.values.get()).isEmpty()) {
				List<Float> list = (List<Float>)this.values.get();
				if (uniformType.isIntStorage()) {
					renderPass.setUniform(this.name, (int)((Float)list.getFirst()).floatValue());
				} else {
					float[] fs = new float[uniformType.getCount()];
					if (list.size() == 1) {
						Arrays.fill(fs, (Float)list.getFirst());
					} else {
						for (int i = 0; i < Math.min(list.size(), uniformType.getCount()); i++) {
							fs[i] = (Float)list.get(i);
						}
					}

					renderPass.setUniform(this.name, fs);
				}
			}
		}
	}
}