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

import com.google.common.collect.Lists;
import com.google.gson.JsonArray;
import com.google.gson.JsonDeserializationContext;
import com.google.gson.JsonDeserializer;
import com.google.gson.JsonElement;
import com.google.gson.JsonParseException;
import java.lang.reflect.Type;
import java.util.List;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.resources.model.BakedModel;
import net.minecraft.client.resources.model.ModelBaker;
import net.minecraft.client.resources.model.ResolvableModel;
import net.minecraft.client.resources.model.WeightedBakedModel;
import net.minecraft.util.random.SimpleWeightedRandomList;
import net.minecraft.world.level.block.state.BlockState;

@Environment(EnvType.CLIENT)
public record MultiVariant(List<Variant> variants) implements UnbakedBlockStateModel {
	public MultiVariant(List<Variant> variants) {
		if (variants.isEmpty()) {
			throw new IllegalArgumentException("Variant list must contain at least one element");
		} else {
			this.variants = variants;
		}
	}

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

	@Override
	public void resolveDependencies(ResolvableModel.Resolver resolver) {
		this.variants.forEach(variant -> resolver.resolve(variant.modelLocation()));
	}

	@Override
	public BakedModel bake(ModelBaker baker) {
		if (this.variants.size() == 1) {
			Variant variant = (Variant)this.variants.getFirst();
			return baker.bake(variant.modelLocation(), variant);
		} else {
			SimpleWeightedRandomList.Builder<BakedModel> builder = SimpleWeightedRandomList.builder();

			for (Variant variant2 : this.variants) {
				BakedModel bakedModel = baker.bake(variant2.modelLocation(), variant2);
				builder.add(bakedModel, variant2.weight());
			}

			return new WeightedBakedModel(builder.build());
		}
	}

	@Environment(EnvType.CLIENT)
	public static class Deserializer implements JsonDeserializer<MultiVariant> {
		public MultiVariant deserialize(JsonElement json, Type type, JsonDeserializationContext context) throws JsonParseException {
			List<Variant> list = Lists.<Variant>newArrayList();
			if (json.isJsonArray()) {
				JsonArray jsonArray = json.getAsJsonArray();
				if (jsonArray.isEmpty()) {
					throw new JsonParseException("Empty variant array");
				}

				for (JsonElement jsonElement : jsonArray) {
					list.add((Variant)context.deserialize(jsonElement, Variant.class));
				}
			} else {
				list.add((Variant)context.deserialize(json, Variant.class));
			}

			return new MultiVariant(list);
		}
	}
}