package net.minecraft.client.renderer.block.model;

import com.google.common.collect.Lists;
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.ObjectArrayList;
import java.util.ArrayList;
import java.util.List;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.renderer.block.model.BlockStateModel.SimpleCachedUnbakedRoot.1;
import net.minecraft.client.renderer.texture.TextureAtlasSprite;
import net.minecraft.client.resources.model.ModelBaker;
import net.minecraft.client.resources.model.ResolvableModel;
import net.minecraft.client.resources.model.WeightedVariants;
import net.minecraft.util.ExtraCodecs;
import net.minecraft.util.RandomSource;
import net.minecraft.util.random.Weighted;
import net.minecraft.util.random.WeightedList;
import net.minecraft.world.level.block.state.BlockState;

@Environment(EnvType.CLIENT)
public interface BlockStateModel {
	void collectParts(RandomSource random, List<BlockModelPart> output);

	default List<BlockModelPart> collectParts(RandomSource random) {
		List<BlockModelPart> list = new ObjectArrayList<>();
		this.collectParts(random, list);
		return list;
	}

	TextureAtlasSprite particleIcon();

	@Environment(EnvType.CLIENT)
	public static class SimpleCachedUnbakedRoot implements BlockStateModel.UnbakedRoot {
		final BlockStateModel.Unbaked contents;
		private final ModelBaker.SharedOperationKey<BlockStateModel> bakingKey = new 1(this);

		public SimpleCachedUnbakedRoot(BlockStateModel.Unbaked contents) {
			this.contents = contents;
		}

		@Override
		public void resolveDependencies(ResolvableModel.Resolver resolver) {
			this.contents.resolveDependencies(resolver);
		}

		@Override
		public BlockStateModel bake(BlockState state, ModelBaker baker) {
			return baker.compute(this.bakingKey);
		}

		@Override
		public Object visualEqualityGroup(BlockState state) {
			return this;
		}
	}

	@Environment(EnvType.CLIENT)
	public interface Unbaked extends ResolvableModel {
		Codec<Weighted<Variant>> ELEMENT_CODEC = RecordCodecBuilder.create(
			instance -> instance.group(Variant.MAP_CODEC.forGetter(Weighted::value), ExtraCodecs.POSITIVE_INT.optionalFieldOf("weight", 1).forGetter(Weighted::weight))
				.apply(instance, Weighted::new)
		);
		Codec<WeightedVariants.Unbaked> HARDCODED_WEIGHTED_CODEC = ExtraCodecs.nonEmptyList(ELEMENT_CODEC.listOf())
			.flatComapMap(
				list -> new WeightedVariants.Unbaked(WeightedList.of(Lists.transform(list, weighted -> weighted.map(SingleVariant.Unbaked::new)))), unbaked -> {
					List<Weighted<BlockStateModel.Unbaked>> list = unbaked.entries().unwrap();
					List<Weighted<Variant>> list2 = new ArrayList(list.size());

					for (Weighted<BlockStateModel.Unbaked> weighted : list) {
						if (!(weighted.value() instanceof SingleVariant.Unbaked unbaked2)) {
							return DataResult.error(() -> "Only single variants are supported");
						}

						list2.add(new Weighted<>(unbaked2.variant(), weighted.weight()));
					}

					return DataResult.success(list2);
				}
			);
		Codec<BlockStateModel.Unbaked> CODEC = Codec.either(HARDCODED_WEIGHTED_CODEC, SingleVariant.Unbaked.CODEC)
			.flatComapMap(either -> either.map(unbaked -> unbaked, unbaked -> unbaked), unbaked -> {
				return switch (unbaked) {
					case SingleVariant.Unbaked unbaked3 -> DataResult.success(Either.right(unbaked3));
					case WeightedVariants.Unbaked unbaked4 -> DataResult.success(Either.left(unbaked4));
					default -> DataResult.error(() -> "Only a single variant or a list of variants are supported");
				};
			});

		BlockStateModel bake(ModelBaker baker);

		default BlockStateModel.UnbakedRoot asRoot() {
			return new BlockStateModel.SimpleCachedUnbakedRoot(this);
		}
	}

	@Environment(EnvType.CLIENT)
	public interface UnbakedRoot extends ResolvableModel {
		BlockStateModel bake(BlockState state, ModelBaker baker);

		Object visualEqualityGroup(BlockState state);
	}
}