package net.minecraft.client.resources.model;

import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.renderer.block.model.BakedQuad;
import net.minecraft.core.Direction;
import net.minecraft.util.RandomSource;
import net.minecraft.world.level.block.state.BlockState;
import org.jetbrains.annotations.Nullable;

@Environment(EnvType.CLIENT)
public class MultiPartBakedModel extends DelegateBakedModel {
	private final List<MultiPartBakedModel.Selector> selectors;
	private final Map<BlockState, BitSet> selectorCache = new Reference2ObjectOpenHashMap<>();

	private static BakedModel getFirstModel(List<MultiPartBakedModel.Selector> list) {
		if (list.isEmpty()) {
			throw new IllegalArgumentException("Model must have at least one selector");
		} else {
			return ((MultiPartBakedModel.Selector)list.getFirst()).model();
		}
	}

	public MultiPartBakedModel(List<MultiPartBakedModel.Selector> selectors) {
		super(getFirstModel(selectors));
		this.selectors = selectors;
	}

	@Override
	public List<BakedQuad> getQuads(@Nullable BlockState state, @Nullable Direction direction, RandomSource random) {
		if (state == null) {
			return Collections.emptyList();
		} else {
			BitSet bitSet = (BitSet)this.selectorCache.get(state);
			if (bitSet == null) {
				bitSet = new BitSet();

				for (int i = 0; i < this.selectors.size(); i++) {
					if (((MultiPartBakedModel.Selector)this.selectors.get(i)).condition.test(state)) {
						bitSet.set(i);
					}
				}

				this.selectorCache.put(state, bitSet);
			}

			List<BakedQuad> list = new ArrayList();
			long l = random.nextLong();

			for (int j = 0; j < bitSet.length(); j++) {
				if (bitSet.get(j)) {
					random.setSeed(l);
					list.addAll(((MultiPartBakedModel.Selector)this.selectors.get(j)).model.getQuads(state, direction, random));
				}
			}

			return list;
		}
	}

	@Environment(EnvType.CLIENT)
	public record Selector(Predicate<BlockState> condition, BakedModel model) {
	}
}