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

import java.util.Map;
import net.minecraft.core.Registry;
import com.google.common.collect.ImmutableMap;
import com.mojang.datafixers.types.DynamicOps;
import net.minecraft.core.BlockPos;
import java.util.Random;
import com.mojang.datafixers.Dynamic;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.entity.ai.behavior.WeightedList;

public class WeightedStateProvider extends BlockStateProvider {
    private final WeightedList<BlockState> weightedList;
    
    private WeightedStateProvider(final WeightedList<BlockState> aod) {
        super(BlockStateProviderType.WEIGHTED_STATE_PROVIDER);
        this.weightedList = aod;
    }
    
    public WeightedStateProvider() {
        this(new WeightedList<BlockState>());
    }
    
    public <T> WeightedStateProvider(final Dynamic<T> dynamic) {
        this(new WeightedList<BlockState>((Dynamic<T>)dynamic.get("entries").orElseEmptyList(), BlockState::deserialize));
    }
    
    public WeightedStateProvider add(final BlockState byj, final int integer) {
        this.weightedList.add(byj, integer);
        return this;
    }
    
    @Override
    public BlockState getState(final Random random, final BlockPos fk) {
        return this.weightedList.getOne(random);
    }
    
    @Override
    public <T> T serialize(final DynamicOps<T> dynamicOps) {
        final ImmutableMap.Builder<T, T> builder3 = ImmutableMap.<T, T>builder();
        builder3.put((T)dynamicOps.createString("type"), (T)dynamicOps.createString(Registry.BLOCKSTATE_PROVIDER_TYPES.getKey(this.type).toString())).put((T)dynamicOps.createString("entries"), this.weightedList.<T>serialize(dynamicOps, byj -> BlockState.<T>serialize(dynamicOps, byj)));
        return (T)new Dynamic((DynamicOps)dynamicOps, dynamicOps.createMap((Map)builder3.build())).getValue();
    }
}
