Revise paste MysticTreeFeature

package com.mysticsbiomes.common.levelgen.feature.tree;

import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.mojang.serialization.Codec;
import com.mysticsbiomes.common.levelgen.feature.config.NewMysticTreeConfiguration;
import net.minecraft.core.BlockPos;
import net.minecraft.core.Direction;
import net.minecraft.tags.BlockTags;
import net.minecraft.util.RandomSource;
import net.minecraft.world.level.LevelAccessor;
import net.minecraft.world.level.LevelSimulatedReader;
import net.minecraft.world.level.WorldGenLevel;
import net.minecraft.world.level.block.LeavesBlock;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.level.block.state.properties.BlockStateProperties;
import net.minecraft.world.level.levelgen.feature.Feature;
import net.minecraft.world.level.levelgen.feature.FeaturePlaceContext;
import net.minecraft.world.level.levelgen.feature.treedecorators.TreeDecorator;
import net.minecraft.world.level.levelgen.structure.BoundingBox;
import net.minecraft.world.level.levelgen.structure.templatesystem.StructureTemplate;
import net.minecraft.world.phys.shapes.BitSetDiscreteVoxelShape;
import net.minecraft.world.phys.shapes.DiscreteVoxelShape;

import java.util.*;
import java.util.function.BiConsumer;

public abstract class NewMysticTreeFeature extends Feature<NewMysticTreeConfiguration> {

public NewMysticTreeFeature(Codec<NewMysticTreeConfiguration> codec) {
        super(codec);
    }

@Override
    public boolean place(FeaturePlaceContext<NewMysticTreeConfiguration> context) {
        final WorldGenLevel level = context.level();
        RandomSource random = context.random();
        BlockPos initialPos = context.origin();
        NewMysticTreeConfiguration config = context.config();

Set<BlockPos> trunkPositions = Sets.newHashSet();
        Set<BlockPos> branchPositions = Sets.newHashSet();
        final Set<BlockPos> foliagePositions = Sets.newHashSet();
        Set<BlockPos> decoratorPositions = Sets.newHashSet();

BiConsumer<BlockPos, BlockState> trunkSetter = setter(level, trunkPositions);
        BiConsumer<BlockPos, BlockState> branchSetter = setter(level, branchPositions);
        BiConsumer<BlockPos, BlockState> foliageSetter = setter(level, foliagePositions);
        BiConsumer<BlockPos, BlockState> decoratorSetter = setter(level, decoratorPositions);

boolean flag = this.doPlace(level, random, initialPos, trunkSetter, branchSetter, foliageSetter, config);
        if (flag && (!branchPositions.isEmpty() || !foliagePositions.isEmpty())) {
            if (!config.decorators.isEmpty()) {
                config.decorators.forEach((decorator) -> decorator.place(new TreeDecorator.Context(level, decoratorSetter, random, branchPositions, foliagePositions, trunkPositions)));
            }
            return BoundingBox.encapsulatingPositions(Iterables.concat(trunkPositions, branchPositions, foliagePositions, decoratorPositions)).map((box) -> {
                StructureTemplate.updateShapeAtEdge(level, 3, updateLeaves(level, box, branchPositions, decoratorPositions, trunkPositions), box.minX(), box.minY(), box.minZ());
                return true;
            }).orElse(false);
        } else {
            return false;
        }
    }

private static BiConsumer<BlockPos, BlockState> setter(LevelAccessor level, Set<BlockPos> set) {
        return (pos, state) -> {
            set.add(pos.immutable());
            level.setBlock(pos, state, 19);
        };
    }

public boolean doPlace(WorldGenLevel level, RandomSource random, BlockPos initialPos, BiConsumer<BlockPos, BlockState> trunkSetter, BiConsumer<BlockPos, BlockState> branchSetter, BiConsumer<BlockPos, BlockState> foliageSetter, NewMysticTreeConfiguration config) {
        int trunkHeight = config.trunkShape.getTrunkHeight(random);
        int minBuildY = Math.min(initialPos.getY(), initialPos.getY());
        int maxBuildY = Math.max(initialPos.getY(), initialPos.getY()) + trunkHeight + 1;

if (minBuildY >= level.getMinBuildHeight() + 1 && maxBuildY <= level.getMaxBuildHeight()) {
            for (int currentY = 0; currentY <= trunkHeight; currentY++) {
                config.trunkShape.placeLog(level, random, initialPos.above(currentY), Direction.Axis.Y, trunkSetter, config);
            }

int foliageRadius = config.foliageShape.radius.sample(random);
            int foliageHeight = config.foliageShape.height.sample(random);
            this.placeFoliage(level, random, initialPos.above(trunkHeight), config, foliageSetter, foliageRadius, foliageHeight);

List<Direction> list = new ArrayList<>();
            for (int count = 0; count < config.branchShape.branchCount.sample(random); count++) {
                Direction branchDirection = Direction.Plane.HORIZONTAL.getRandomDirection(random);
                list.add(branchDirection);

for (Direction directions : list) {
                    if (branchDirection != directions) {
                        int branchYStart = trunkHeight / 2 - config.branchShape.branchYStart.sample(random);
                        this.generateBranch(level, random, initialPos.above(branchYStart), branchDirection, branchSetter, foliageSetter, config);
                    }
                }
            }
            return true;
        } else {
            return false;
        }
    }
    
    private int getMaxFreeTreeHeight(LevelSimulatedReader level, RandomSource random, BlockPos initialPos, int height, NewMysticTreeConfiguration config) {
        BlockPos.MutableBlockPos mutablePos = new BlockPos.MutableBlockPos();

for (int y = 0; y <= height + 1; y++) {
            int radius = config.foliageShape.foliageRadius(random);

for (int xOffset = -radius; xOffset <= radius; xOffset++) {
                for (int zOffset = -radius; zOffset <= radius; zOffset++) {
                    mutablePos.setWithOffset(initialPos, xOffset, y, zOffset);

boolean flag = !config.trunkShape.isFree(level, mutablePos);
                    if (flag) {
                        return y - 2;
                    }
                }
            }
        }
        return height;
    }

private static DiscreteVoxelShape updateLeaves(LevelAccessor level, BoundingBox box, Set<BlockPos> logs, Set<BlockPos> leaves, Set<BlockPos> persistentLeaves) {
        DiscreteVoxelShape shape = new BitSetDiscreteVoxelShape(box.getXSpan(), box.getYSpan(), box.getZSpan());
        List<Set<BlockPos>> list = Lists.newArrayList();

for (int i = 0; i < 7; i++) {
            list.add(Sets.newHashSet());
        }

for (BlockPos pos : Lists.newArrayList(Sets.union(leaves, persistentLeaves))) {
            if (box.isInside(pos)) {
                shape.fill(pos.getX() - box.minX(), pos.getY() - box.minY(), pos.getZ() - box.minZ());
            }
        }

BlockPos.MutableBlockPos mutablePos = new BlockPos.MutableBlockPos();
        int currentDistance = 0;
        list.getFirst().addAll(logs);

while (true) {
            while (currentDistance >= 7 || !list.get(currentDistance).isEmpty()) {
                if (currentDistance >= 7) return shape;

Iterator<BlockPos> iterator = list.get(currentDistance).iterator();
                BlockPos currentPos = iterator.next();
                iterator.remove();

if (box.isInside(currentPos)) {
                    if (currentDistance != 0) {
                        BlockState state = level.getBlockState(currentPos);
                        level.setBlock(currentPos, state.setValue(BlockStateProperties.DISTANCE, 1), 19);
                    }

shape.fill(currentPos.getX() - box.minX(), currentPos.getY() - box.minY(), currentPos.getZ() - box.minZ());

for (Direction direction : Direction.values()) {
                        mutablePos.setWithOffset(currentPos, direction);
                        if (box.isInside(mutablePos)) {
                            int dx = mutablePos.getX() - box.minX();
                            int dy = mutablePos.getY() - box.minY();
                            int dz = mutablePos.getZ() - box.minZ();

if (!shape.isFull(dx, dy, dz)) {
                                OptionalInt distance = LeavesBlock.getOptionalDistanceAt(level.getBlockState(mutablePos));
                                if (distance.isPresent()) {
                                    int nextDistance = Math.min(distance.getAsInt(), currentDistance + 1);
                                    if (nextDistance < 7) {
                                        list.get(nextDistance).add(mutablePos.immutable());
                                        currentDistance = Math.min(currentDistance, nextDistance);
                                    }
                                }
                            }
                        }
                    }
                }
            }
            ++currentDistance;
        }
    }

public static boolean validTreePos(LevelSimulatedReader level, BlockPos pos) {
        return level.isStateAtPosition(pos, state -> state.isAir() || state.is(BlockTags.REPLACEABLE_BY_TREES));
    }

public static boolean isAirOrLeaves(LevelSimulatedReader level, BlockPos pos) {
        return level.isStateAtPosition(pos, state -> state.isAir() || state.is(BlockTags.LEAVES));
    }

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

protected void generateBranch(LevelAccessor level, RandomSource random, BlockPos startPos, Direction branchDirection, BiConsumer<BlockPos, BlockState> branchSetter, BiConsumer<BlockPos, BlockState> foliageSetter, NewMysticTreeConfiguration config) {
        BlockPos.MutableBlockPos mutablePos = startPos.mutable();

int branchLength = config.branchShape.branchLength.sample(random);
        int extendAmount = config.branchShape.extendAmount.sample(random);
        if (extendAmount > 0) {
            for (int i = 0; i < extendAmount; i++) {
                mutablePos.move(branchDirection);
                config.trunkShape.placeLog(level, random, mutablePos, branchDirection.getAxis(), branchSetter, config);
            }
        }

int steps = 0;
        while (steps <= branchLength) {
            int forwardInterval = 1 + Math.round((float) steps / branchLength * (config.branchShape.forwardInterval + 1));
            int upwardInterval = 1 + Math.round((float) steps / branchLength * (config.branchShape.upwardInterval + 1));

for (int i = 0; i < forwardInterval && steps < branchLength; i++, steps++) {
                mutablePos.move(branchDirection);
                config.trunkShape.placeLog(level, random, mutablePos, branchDirection.getAxis(), branchSetter, config);
            }

for (int i = 0; i < upwardInterval && steps < branchLength; i++, steps++) {
                mutablePos.move(Direction.UP);
                config.trunkShape.placeLog(level, random, mutablePos, Direction.Axis.Y, branchSetter, config);
            }

if (steps == branchLength) {
                this.placeFoliage(level, random, mutablePos, config, foliageSetter, config.foliageShape.radius.sample(random), config.foliageShape.height.sample(random));
                return;
            }
        }
    }

public abstract void placeFoliage(LevelAccessor level, RandomSource random, BlockPos pos, NewMysticTreeConfiguration config, BiConsumer<BlockPos, BlockState> foliageSetter, int foliageRadius, int foliageHeight);

}