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.block.model.TextureSlots;
import net.minecraft.client.renderer.special.SpecialModelRenderer;
import net.minecraft.client.renderer.special.SpecialModelRenderers;
import net.minecraft.client.resources.model.ModelBaker;
import net.minecraft.client.resources.model.ResolvableModel;
import net.minecraft.client.resources.model.ResolvedModel;
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;
import org.joml.Vector3f;

@Environment(EnvType.CLIENT)
public class SpecialModelWrapper<T> implements ItemModel {
	private static final Vector3f[] EXTENTS = new Vector3f[]{
		new Vector3f(0.0F, 0.0F, 0.0F),
		new Vector3f(0.0F, 0.0F, 1.0F),
		new Vector3f(0.0F, 1.0F, 1.0F),
		new Vector3f(0.0F, 1.0F, 0.0F),
		new Vector3f(1.0F, 1.0F, 0.0F),
		new Vector3f(1.0F, 1.0F, 1.0F),
		new Vector3f(1.0F, 0.0F, 1.0F),
		new Vector3f(1.0F, 0.0F, 0.0F)
	};
	private final SpecialModelRenderer<T> specialRenderer;
	private final ModelRenderProperties properties;

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

	@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.setExtents(() -> EXTENTS);
		layerRenderState.setupSpecialModel(this.specialRenderer, this.specialRenderer.extractArgument(stack));
		this.properties.applyToLayer(layerRenderState, displayContext);
	}

	@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.markDependency(this.base);
		}

		@Override
		public ItemModel bake(ItemModel.BakingContext context) {
			SpecialModelRenderer<?> specialModelRenderer = this.specialModel.bake(context.entityModelSet());
			if (specialModelRenderer == null) {
				return context.missingItemModel();
			} else {
				ModelRenderProperties modelRenderProperties = this.getProperties(context);
				return new SpecialModelWrapper<>(specialModelRenderer, modelRenderProperties);
			}
		}

		private ModelRenderProperties getProperties(ItemModel.BakingContext context) {
			ModelBaker modelBaker = context.blockModelBaker();
			ResolvedModel resolvedModel = modelBaker.getModel(this.base);
			TextureSlots textureSlots = resolvedModel.getTopTextureSlots();
			return ModelRenderProperties.fromResolvedModel(modelBaker, resolvedModel, textureSlots);
		}

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