package net.minecraft.client.renderer.item;

import com.mojang.serialization.MapCodec;
import com.mojang.serialization.codecs.RecordCodecBuilder;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.multiplayer.ClientLevel;
import net.minecraft.client.renderer.special.SpecialModelRenderer;
import net.minecraft.client.renderer.special.SpecialModelRenderers;
import net.minecraft.client.resources.model.BakedModel;
import net.minecraft.client.resources.model.ResolvableModel;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.item.ItemDisplayContext;
import net.minecraft.world.item.ItemStack;
import org.jetbrains.annotations.Nullable;

@Environment(EnvType.CLIENT)
public class SpecialModelWrapper<T> implements ItemModel {
	private final SpecialModelRenderer<T> specialRenderer;
	private final BakedModel baseModel;

	public SpecialModelWrapper(SpecialModelRenderer<T> specialRenderer, BakedModel baseModel) {
		this.specialRenderer = specialRenderer;
		this.baseModel = baseModel;
	}

	@Override
	public void update(
		ItemStackRenderState renderState,
		ItemStack stack,
		ItemModelResolver itemModelResolver,
		ItemDisplayContext displayContext,
		@Nullable ClientLevel level,
		@Nullable LivingEntity entity,
		int seed
	) {
		ItemStackRenderState.LayerRenderState layerRenderState = renderState.newLayer();
		if (stack.hasFoil()) {
			layerRenderState.setFoilType(ItemStackRenderState.FoilType.STANDARD);
		}

		layerRenderState.setupSpecialModel(this.specialRenderer, this.specialRenderer.extractArgument(stack), this.baseModel);
	}

	@Environment(EnvType.CLIENT)
	public record Unbaked(ResourceLocation base, SpecialModelRenderer.Unbaked specialModel) implements ItemModel.Unbaked {
		public static final MapCodec<SpecialModelWrapper.Unbaked> MAP_CODEC = RecordCodecBuilder.mapCodec(
			instance -> instance.group(
					ResourceLocation.CODEC.fieldOf("base").forGetter(SpecialModelWrapper.Unbaked::base),
					SpecialModelRenderers.CODEC.fieldOf("model").forGetter(SpecialModelWrapper.Unbaked::specialModel)
				)
				.apply(instance, SpecialModelWrapper.Unbaked::new)
		);

		@Override
		public void resolveDependencies(ResolvableModel.Resolver resolver) {
			resolver.resolve(this.base);
		}

		@Override
		public ItemModel bake(ItemModel.BakingContext context) {
			BakedModel bakedModel = context.bake(this.base);
			SpecialModelRenderer<?> specialModelRenderer = this.specialModel.bake(context.entityModelSet());
			return (ItemModel)(specialModelRenderer == null ? context.missingItemModel() : new SpecialModelWrapper<>(specialModelRenderer, bakedModel));
		}

		@Override
		public MapCodec<SpecialModelWrapper.Unbaked> type() {
			return MAP_CODEC;
		}
	}
}