package net.minecraft.client.model.geom.builders;

import com.google.common.collect.Maps;
import it.unimi.dsi.fastutil.objects.Object2ObjectArrayMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.Map.Entry;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.model.geom.ModelPart;
import net.minecraft.client.model.geom.PartPose;

@Environment(EnvType.CLIENT)
public class PartDefinition {
	private final List<CubeDefinition> cubes;
	private final PartPose partPose;
	private final Map<String, PartDefinition> children = Maps.<String, PartDefinition>newHashMap();

	PartDefinition(List<CubeDefinition> cubes, PartPose partPose) {
		this.cubes = cubes;
		this.partPose = partPose;
	}

	public PartDefinition addOrReplaceChild(String name, CubeListBuilder cubes, PartPose partPose) {
		PartDefinition partDefinition = new PartDefinition(cubes.getCubes(), partPose);
		return this.addOrReplaceChild(name, partDefinition);
	}

	public PartDefinition addOrReplaceChild(String name, PartDefinition chuld) {
		PartDefinition partDefinition = (PartDefinition)this.children.put(name, chuld);
		if (partDefinition != null) {
			chuld.children.putAll(partDefinition.children);
		}

		return chuld;
	}

	public PartDefinition clearChild(String name) {
		PartDefinition partDefinition = (PartDefinition)this.children.get(name);
		if (partDefinition == null) {
			throw new IllegalArgumentException("No child with name: " + name);
		} else {
			return this.addOrReplaceChild(name, CubeListBuilder.create(), partDefinition.partPose);
		}
	}

	public ModelPart bake(int texWidth, int texHeight) {
		Object2ObjectArrayMap<String, ModelPart> object2ObjectArrayMap = (Object2ObjectArrayMap<String, ModelPart>)this.children
			.entrySet()
			.stream()
			.collect(
				Collectors.toMap(
					Entry::getKey, entry -> ((PartDefinition)entry.getValue()).bake(texWidth, texHeight), (modelPartx, modelPart2) -> modelPartx, Object2ObjectArrayMap::new
				)
			);
		List<ModelPart.Cube> list = this.cubes.stream().map(cubeDefinition -> cubeDefinition.bake(texWidth, texHeight)).toList();
		ModelPart modelPart = new ModelPart(list, object2ObjectArrayMap);
		modelPart.setInitialPose(this.partPose);
		modelPart.loadPose(this.partPose);
		return modelPart;
	}

	public PartDefinition getChild(String name) {
		return (PartDefinition)this.children.get(name);
	}

	public Set<Entry<String, PartDefinition>> getChildren() {
		return this.children.entrySet();
	}

	public PartDefinition transformed(UnaryOperator<PartPose> transformer) {
		PartDefinition partDefinition = new PartDefinition(this.cubes, (PartPose)transformer.apply(this.partPose));
		partDefinition.children.putAll(this.children);
		return partDefinition;
	}
}