package com.r3944realms.leashedplayer.client.renderer.item;

import com.mojang.serialization.Codec;
import com.mojang.serialization.MapCodec;
import com.mojang.serialization.codecs.RecordCodecBuilder;
import com.r3944realms.leashedplayer.client.renderer.item.properties.conditionalRange.ConditionalRangeItemModelProperties;
import com.r3944realms.leashedplayer.client.renderer.item.properties.conditionalRange.ConditionalRangeItemModelProperty;
import net.minecraft.client.multiplayer.ClientLevel;
import net.minecraft.client.renderer.item.ItemModel;
import net.minecraft.client.renderer.item.ItemModelResolver;
import net.minecraft.client.renderer.item.ItemModels;
import net.minecraft.client.renderer.item.ItemStackRenderState;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.item.ItemDisplayContext;
import net.minecraft.world.item.ItemStack;
import net.neoforged.api.distmarker.Dist;
import net.neoforged.api.distmarker.OnlyIn;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

@OnlyIn(Dist.CLIENT)
public class ConditionalRangeItemModel implements ItemModel {
    private final ConditionalRangeItemModelProperty property;
    private final float scale;
    private final float[] onTrueThresholds;
    private final float[] onFalseThresholds;
    private final ItemModel[] onTrueModels;
    private final ItemModel[] onFalseModels;
    private final ItemModel fallback;
    public ConditionalRangeItemModel(ConditionalRangeItemModelProperty property, float scale, float[] onTrueThresholds, float[] onFalseThresholds, ItemModel[] onTrueModels, ItemModel[] onFalseModels, ItemModel fallback) {
        this.property = property;
        this.scale = scale;
        this.onTrueThresholds = onTrueThresholds;
        this.onFalseThresholds = onFalseThresholds;
        this.onTrueModels = onTrueModels;
        this.onFalseModels = onFalseModels;
        this.fallback = fallback;
    }
    private static int lastIndexLessOrEqual(float[] thresholds, float value) {
        if (thresholds.length < 16) {
            for (int k = 0; k < thresholds.length; k++) {
                if (thresholds[k] > value) {
                    return k - 1;
                }
            }

            return thresholds.length - 1;
        } else {
            int i = Arrays.binarySearch(thresholds, value);
            if (i < 0) {
                int j = ~i;
                return j - 1;
            } else {
                return i;
            }
        }
    }
    @Override
    public void update(@NotNull ItemStackRenderState renderState, @NotNull ItemStack stack, @NotNull ItemModelResolver itemModelResolver, @NotNull ItemDisplayContext displayContext, @Nullable ClientLevel level, @Nullable LivingEntity entity, int seed) {
        ConditionalRangeItemModelProperty.ConditionalRangeData conditionalRangeData = this.property.get(stack, level, entity, seed, displayContext);
        float f = conditionalRangeData.floatValue() * this.scale;
        boolean b = conditionalRangeData.boolValue();
        ItemModel itemModel;
        if(Float.isNaN(f)) {
            itemModel = this.fallback; ;
        } else {
            int i;
            if (b) {
                i = lastIndexLessOrEqual(this.onTrueThresholds, f);
            }
            else {
                i = lastIndexLessOrEqual(this.onFalseThresholds, f);
            }
            itemModel = i == -1 ? this.fallback : b ? onTrueModels[i] : onFalseModels[i];

        }
        itemModel.update(renderState, stack, itemModelResolver, displayContext, level, entity, seed);
    }
    public record Entry(float threshold, ItemModel.Unbaked model) {
        public static final Codec<Entry> CODEC = RecordCodecBuilder.create(
                instance -> instance.group(
                        Codec.FLOAT.fieldOf("threshold").forGetter(Entry::threshold),
                        ItemModels.CODEC.fieldOf("model").forGetter(Entry::model)
                ).apply(instance, Entry::new)
        );
        public static final Comparator<Entry> BY_THRESHOLD = Comparator.comparingDouble(Entry::threshold);
    }
    public record Unbaked(
            ConditionalRangeItemModelProperty property, float scale, List<Entry> trueModels, List<Entry> falseModels, Optional<ItemModel.Unbaked> fallback
    ) implements ItemModel.Unbaked {
        public static final MapCodec<Unbaked> MAP_CODEC = RecordCodecBuilder.mapCodec(
          instance -> instance.group(
                  ConditionalRangeItemModelProperties.MAP_CODEC.forGetter(Unbaked::property),
                  Codec.FLOAT.optionalFieldOf("scale", 1.0F).forGetter(Unbaked::scale),
                  Entry.CODEC.listOf().fieldOf("trueModels").forGetter(Unbaked::trueModels),
                  Entry.CODEC.listOf().fieldOf("falseModels").forGetter(Unbaked::falseModels),
                  ItemModels.CODEC.optionalFieldOf("fallback").forGetter(Unbaked::fallback)
          ).apply(instance, Unbaked::new)
        );
        @Override
        public @NotNull MapCodec<? extends ItemModel.Unbaked> type() {
            return MAP_CODEC;
        }

        @Override
        public @NotNull ItemModel bake(@NotNull BakingContext context) {
            float[] floatArrA = new float[this.trueModels.size()];
            float[] floatArrB = new float[this.falseModels.size()];
            ItemModel[] itemModelArrA = new ItemModel[this.trueModels.size()];
            ItemModel[] itemModelArrB = new ItemModel[this.falseModels.size()];
            List<Entry> trueModelList = new ArrayList<>(this.trueModels);
            List<Entry> falseModelList = new ArrayList<>(this.falseModels);
            trueModelList.sort(Entry.BY_THRESHOLD);
            falseModelList.sort(Entry.BY_THRESHOLD);
            for (int i = 0; i < trueModelList.size(); i++) {
                Entry rangeselectitemmodel$entry = trueModelList.get(i);
                floatArrA[i] = rangeselectitemmodel$entry.threshold;
                itemModelArrA[i] = rangeselectitemmodel$entry.model.bake(context);
            }
            for (int i = 0; i < falseModels.size(); i++) {
                Entry rangeselectitemmodel$entry = falseModelList.get(i);
                floatArrB[i] = rangeselectitemmodel$entry.threshold;
                itemModelArrB[i] = rangeselectitemmodel$entry.model.bake(context);
            }
            ItemModel itemmodel = this.fallback.map(p_387030_ -> p_387030_.bake(context)).orElse(context.missingItemModel());
            return new ConditionalRangeItemModel(this.property, this.scale, floatArrA, floatArrB, itemModelArrA, itemModelArrB, itemmodel);
        }

        @Override
        public void resolveDependencies(@NotNull Resolver resolver) {
            this.fallback.ifPresent(itemModel -> itemModel.resolveDependencies(resolver));
            this.trueModels.forEach(entry -> entry.model.resolveDependencies(resolver));
            this.falseModels.forEach(entry -> entry.model.resolveDependencies(resolver));
        }
    }

}
