package net.minecraft.world.level.chunk;

import java.util.AbstractList;
import net.minecraft.nbt.ListTag;
import net.minecraft.network.FriendlyByteBuf;
import javax.annotation.Nullable;
import java.util.function.Predicate;
import net.minecraft.nbt.CompoundTag;
import java.util.function.Function;
import net.minecraft.util.CrudeIncrementalIntIdentityHashBiMap;
import net.minecraft.core.IdMapper;

public class HashMapPalette<T> implements Palette<T> {
    private final IdMapper<T> registry;
    private final CrudeIncrementalIntIdentityHashBiMap<T> values;
    private final PaletteResize<T> resizeHandler;
    private final Function<CompoundTag, T> reader;
    private final Function<T, CompoundTag> writer;
    private final int bits;
    
    public HashMapPalette(final IdMapper<T> ge, final int integer, final PaletteResize<T> chp, final Function<CompoundTag, T> function4, final Function<T, CompoundTag> function5) {
        this.registry = ge;
        this.bits = integer;
        this.resizeHandler = chp;
        this.reader = function4;
        this.writer = function5;
        this.values = new CrudeIncrementalIntIdentityHashBiMap<T>(1 << integer);
    }
    
    @Override
    public int idFor(final T object) {
        int integer3 = this.values.getId(object);
        if (integer3 == -1) {
            integer3 = this.values.add(object);
            if (integer3 >= 1 << this.bits) {
                integer3 = this.resizeHandler.onResize(this.bits + 1, object);
            }
        }
        return integer3;
    }
    
    @Override
    public boolean maybeHas(final Predicate<T> predicate) {
        for (int integer3 = 0; integer3 < this.getSize(); ++integer3) {
            if (predicate.test(this.values.byId(integer3))) {
                return true;
            }
        }
        return false;
    }
    
    @Nullable
    @Override
    public T valueFor(final int integer) {
        return this.values.byId(integer);
    }
    
    @Override
    public void read(final FriendlyByteBuf mg) {
        this.values.clear();
        for (int integer3 = mg.readVarInt(), integer4 = 0; integer4 < integer3; ++integer4) {
            this.values.add(this.registry.byId(mg.readVarInt()));
        }
    }
    
    @Override
    public void write(final FriendlyByteBuf mg) {
        final int integer3 = this.getSize();
        mg.writeVarInt(integer3);
        for (int integer4 = 0; integer4 < integer3; ++integer4) {
            mg.writeVarInt(this.registry.getId(this.values.byId(integer4)));
        }
    }
    
    @Override
    public int getSerializedSize() {
        int integer2 = FriendlyByteBuf.getVarIntSize(this.getSize());
        for (int integer3 = 0; integer3 < this.getSize(); ++integer3) {
            integer2 += FriendlyByteBuf.getVarIntSize(this.registry.getId(this.values.byId(integer3)));
        }
        return integer2;
    }
    
    public int getSize() {
        return this.values.size();
    }
    
    @Override
    public void read(final ListTag lk) {
        this.values.clear();
        for (int integer3 = 0; integer3 < lk.size(); ++integer3) {
            this.values.add(this.reader.apply(lk.getCompound(integer3)));
        }
    }
    
    public void write(final ListTag lk) {
        for (int integer3 = 0; integer3 < this.getSize(); ++integer3) {
            ((AbstractList<CompoundTag>)lk).add(this.writer.apply(this.values.byId(integer3)));
        }
    }
}
