package net.minecraft.client.renderer.item;

import java.util.function.Function;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.Minecraft;
import net.minecraft.client.multiplayer.ClientLevel;
import net.minecraft.client.resources.model.ModelManager;
import net.minecraft.core.component.DataComponents;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.entity.Entity;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.item.BlockItem;
import net.minecraft.world.item.ItemDisplayContext;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.item.component.ResolvableProfile;
import net.minecraft.world.level.Level;
import net.minecraft.world.level.block.AbstractSkullBlock;
import org.jetbrains.annotations.Nullable;

@Environment(EnvType.CLIENT)
public class ItemModelResolver {
	private final Function<ResourceLocation, ItemModel> modelGetter;
	private final Function<ResourceLocation, ClientItem.Properties> clientProperties;

	public ItemModelResolver(ModelManager modelManager) {
		this.modelGetter = modelManager::getItemModel;
		this.clientProperties = modelManager::getItemProperties;
	}

	public void updateForLiving(ItemStackRenderState renderState, ItemStack stack, ItemDisplayContext displayContext, LivingEntity entity) {
		this.updateForTopItem(renderState, stack, displayContext, entity.level(), entity, entity.getId() + displayContext.ordinal());
	}

	public void updateForNonLiving(ItemStackRenderState renderState, ItemStack stack, ItemDisplayContext displayContext, Entity entity) {
		this.updateForTopItem(renderState, stack, displayContext, entity.level(), null, entity.getId());
	}

	public void updateForTopItem(
		ItemStackRenderState renderState, ItemStack stack, ItemDisplayContext displayContext, @Nullable Level level, @Nullable LivingEntity entity, int seed
	) {
		renderState.clear();
		if (!stack.isEmpty()) {
			renderState.displayContext = displayContext;
			this.appendItemLayers(renderState, stack, displayContext, level, entity, seed);
		}
	}

	private static void fixupSkullProfile(ItemStack stack) {
		if (stack.getItem() instanceof BlockItem blockItem && blockItem.getBlock() instanceof AbstractSkullBlock) {
			ResolvableProfile resolvableProfile = stack.get(DataComponents.PROFILE);
			if (resolvableProfile != null && !resolvableProfile.isResolved()) {
				stack.remove(DataComponents.PROFILE);
				resolvableProfile.resolve().thenAcceptAsync(resolvableProfilex -> stack.set(DataComponents.PROFILE, resolvableProfilex), Minecraft.getInstance());
			}
		}
	}

	public void appendItemLayers(
		ItemStackRenderState renderState, ItemStack stack, ItemDisplayContext displayContext, @Nullable Level level, @Nullable LivingEntity entity, int seed
	) {
		fixupSkullProfile(stack);
		ResourceLocation resourceLocation = stack.get(DataComponents.ITEM_MODEL);
		if (resourceLocation != null) {
			((ItemModel)this.modelGetter.apply(resourceLocation))
				.update(renderState, stack, this, displayContext, level instanceof ClientLevel clientLevel ? clientLevel : null, entity, seed);
		}
	}

	public boolean shouldPlaySwapAnimation(ItemStack stack) {
		ResourceLocation resourceLocation = stack.get(DataComponents.ITEM_MODEL);
		return resourceLocation == null ? true : ((ClientItem.Properties)this.clientProperties.apply(resourceLocation)).handAnimationOnSwap();
	}
}