package net.minecraft.world.level.levelgen.feature.stateproviders;

import java.util.function.Function;
import net.minecraft.core.BlockPos;
import java.util.Random;
import com.mojang.serialization.DataResult;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.entity.ai.behavior.WeightedList;
import com.mojang.serialization.Codec;

public class WeightedStateProvider extends BlockStateProvider {
    public static final Codec<WeightedStateProvider> CODEC;
    private final WeightedList<BlockState> weightedList;
    
    private static DataResult<WeightedStateProvider> create(final WeightedList<BlockState> atb) {
        if (atb.isEmpty()) {
            return (DataResult<WeightedStateProvider>)DataResult.error("WeightedStateProvider with no states");
        }
        return (DataResult<WeightedStateProvider>)DataResult.success(new WeightedStateProvider(atb));
    }
    
    private WeightedStateProvider(final WeightedList<BlockState> atb) {
        this.weightedList = atb;
    }
    
    @Override
    protected BlockStateProviderType<?> type() {
        return BlockStateProviderType.WEIGHTED_STATE_PROVIDER;
    }
    
    public WeightedStateProvider() {
        this(new WeightedList<BlockState>());
    }
    
    public WeightedStateProvider add(final BlockState cfj, final int integer) {
        this.weightedList.add(cfj, integer);
        return this;
    }
    
    @Override
    public BlockState getState(final Random random, final BlockPos fu) {
        return this.weightedList.getOne(random);
    }
    
    static {
        CODEC = WeightedList.<BlockState>codec(BlockState.CODEC).comapFlatMap((Function)WeightedStateProvider::create, cpu -> cpu.weightedList).fieldOf("entries").codec();
    }
}
