package net.minecraft.client.model;

import com.mojang.blaze3d.vertex.PoseStack;
import com.mojang.blaze3d.vertex.VertexConsumer;
import java.util.List;
import java.util.Optional;
import java.util.function.Function;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.animation.AnimationDefinition;
import net.minecraft.client.animation.KeyframeAnimations;
import net.minecraft.client.model.geom.ModelPart;
import net.minecraft.client.renderer.RenderType;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.entity.AnimationState;
import org.joml.Vector3f;

@Environment(EnvType.CLIENT)
public abstract class Model {
	private static final Vector3f ANIMATION_VECTOR_CACHE = new Vector3f();
	protected final ModelPart root;
	protected final Function<ResourceLocation, RenderType> renderType;
	private final List<ModelPart> allParts;

	public Model(ModelPart root, Function<ResourceLocation, RenderType> renderType) {
		this.root = root;
		this.renderType = renderType;
		this.allParts = root.getAllParts().toList();
	}

	public final RenderType renderType(ResourceLocation location) {
		return (RenderType)this.renderType.apply(location);
	}

	public final void renderToBuffer(PoseStack poseStack, VertexConsumer buffer, int packedLight, int packedOverlay, int color) {
		this.root().render(poseStack, buffer, packedLight, packedOverlay, color);
	}

	public final void renderToBuffer(PoseStack poseStack, VertexConsumer buffer, int packedLight, int packedOverlay) {
		this.renderToBuffer(poseStack, buffer, packedLight, packedOverlay, -1);
	}

	public final ModelPart root() {
		return this.root;
	}

	public Optional<ModelPart> getAnyDescendantWithName(String name) {
		return name.equals("root")
			? Optional.of(this.root())
			: this.root().getAllParts().filter(modelPart -> modelPart.hasChild(name)).findFirst().map(modelPart -> modelPart.getChild(name));
	}

	public final List<ModelPart> allParts() {
		return this.allParts;
	}

	public final void resetPose() {
		for (ModelPart modelPart : this.allParts) {
			modelPart.resetPose();
		}
	}

	protected void animate(AnimationState state, AnimationDefinition definition, float ageInTicks) {
		this.animate(state, definition, ageInTicks, 1.0F);
	}

	protected void animateWalk(AnimationDefinition definition, float walkAnimationPos, float walkAnimationSpeed, float timeMultiplier, float speedMultiplier) {
		long l = (long)(walkAnimationPos * 50.0F * timeMultiplier);
		float f = Math.min(walkAnimationSpeed * speedMultiplier, 1.0F);
		KeyframeAnimations.animate(this, definition, l, f, ANIMATION_VECTOR_CACHE);
	}

	protected void animate(AnimationState state, AnimationDefinition definition, float ageInTicks, float speed) {
		state.ifStarted(
			animationState -> KeyframeAnimations.animate(
				this, definition, (long)((float)animationState.getTimeInMillis(ageInTicks) * speed), 1.0F, ANIMATION_VECTOR_CACHE
			)
		);
	}

	protected void applyStatic(AnimationDefinition animationDefinition) {
		KeyframeAnimations.animate(this, animationDefinition, 0L, 1.0F, ANIMATION_VECTOR_CACHE);
	}

	@Environment(EnvType.CLIENT)
	public static class Simple extends Model {
		public Simple(ModelPart root, Function<ResourceLocation, RenderType> renderType) {
			super(root, renderType);
		}
	}
}