package com.teammetallurgy.atum.world.gen.feature;

import com.google.common.collect.Sets;
import com.mojang.serialization.Codec;
import com.teammetallurgy.atum.blocks.wood.DeadwoodBranchBlock;
import com.teammetallurgy.atum.blocks.wood.DeadwoodLogBlock;
import com.teammetallurgy.atum.init.AtumBlocks;
import net.minecraft.core.BlockPos;
import net.minecraft.core.Direction;
import net.minecraft.core.Vec3i;
import net.minecraft.util.RandomSource;
import net.minecraft.world.level.LevelSimulatedRW;
import net.minecraft.world.level.WorldGenLevel;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.level.levelgen.feature.Feature;
import net.minecraft.world.level.levelgen.feature.FeaturePlaceContext;
import net.minecraft.world.level.levelgen.feature.TreeFeature;
import net.minecraft.world.level.levelgen.feature.configurations.NoneFeatureConfiguration;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;

import javax.annotation.Nonnull;
import java.util.*;

public class DeadwoodFeature extends Feature<NoneFeatureConfiguration> {

    public DeadwoodFeature(Codec<NoneFeatureConfiguration> config) {
        super(config);
    }

    @Override
    public boolean place(@Nonnull FeaturePlaceContext<NoneFeatureConfiguration> placeContext) {
        WorldGenLevel genLevel = placeContext.level();
        BlockPos pos = placeContext.origin();
        RandomSource random = placeContext.random();
        BlockState log = AtumBlocks.DEADWOOD_LOG.get().defaultBlockState().setValue(DeadwoodLogBlock.HAS_SCARAB, true);
        BlockState branch = AtumBlocks.DEADWOOD_BRANCH.get().defaultBlockState();

        Set<BlockPos> logs = Sets.newHashSet();
        int baseHeight = random.nextInt(5) + 3;
        boolean generate = true;

        if (pos.getY() >= 0 && pos.getY() + baseHeight + 1 <= genLevel.getMaxBuildHeight()) {
            for (int y = pos.getY(); y <= pos.getY() + 1 + baseHeight; ++y) {
                int k = 1;

                if (y == pos.getY()) {
                    k = 0;
                }

                if (y >= pos.getY() + 1 + baseHeight - 2) {
                    k = 2;
                }
                BlockPos.MutableBlockPos mutable = new BlockPos.MutableBlockPos();

                for (int x = pos.getX() - k; x <= pos.getX() + k && generate; ++x) {
                    for (int z = pos.getZ() - k; z <= pos.getZ() + k && generate; ++z) {
                        if (y >= 0 && y < genLevel.getMaxBuildHeight()) {
                            BlockPos checkPos = mutable.set(x, y, z);
                            if (!TreeFeature.isAirOrLeaves(genLevel, checkPos)) {
                                generate = false;
                            }
                        } else {
                            generate = false;
                        }
                    }
                }
            }

            if (!generate) {
                return false;
            } else {
                BlockPos down = pos.below();
                BlockState state = genLevel.getBlockState(down);
                boolean isSoil = state.getBlock() == AtumBlocks.STRANGE_SAND.get();

                if (genLevel.isAreaLoaded(pos, 16)) {
                    if (isSoil && pos.getY() < genLevel.getMaxBuildHeight() - baseHeight - 1) {
                        for (int height = 0; height < baseHeight; ++height) {
                            BlockPos upN = pos.above(height);

                            if (TreeFeature.isAirOrLeaves(genLevel, upN)) {
                                genLevel.setBlock(pos.above(height), log, 19);
                                if (height > 1) {
                                    logs.add(pos.above(height));
                                }
                            }
                        }
                        buildBranches(genLevel, genLevel, logs, random, log, branch);
                        return true;
                    } else {
                        return false;
                    }
                } else {
                    return false;
                }
            }
        }
        return false;
    }

    public void buildBranches(WorldGenLevel genLevel, LevelSimulatedRW simulatedRW, Set<BlockPos> logs, RandomSource random, BlockState log, BlockState branch) {
        // Keep track of branches to be generated using a Queue
        Queue<Pair<BlockPos, Integer>> queue = new LinkedList<>();
        List<BlockPos> placedBranches = new ArrayList<>();

        for (BlockPos pos : logs) {
            queue.add(new ImmutablePair<>(pos, 0));
        }
        if (queue.peek() != null) {
            BlockPos baseLog = queue.peek().getLeft();
            int count = 0;

            Set<Direction> choosenFacings = new HashSet<>();
            Map<Direction, Integer> balancingCount = new HashMap<>();
            for (Direction facing : Direction.values()) {
                balancingCount.put(facing, 0);
            }

            while (!queue.isEmpty()) {
                Pair<BlockPos, Integer> pair = queue.poll();
                if (pair.getLeft() == null || pair.getRight() == null) {
                    continue;
                }
                BlockPos pos = pair.getLeft();
                int branchLength = pair.getRight();

                choosenFacings.clear();
                for (Direction facing : Direction.values()) {

                    // Prevent a split from splitting in both directions
                    if (choosenFacings.contains(facing.getOpposite())) {
                        continue;
                    }

                    // Prevent branches 3 blocks in the same direction
                    BlockState curr = genLevel.getBlockState(pos);
                    if (curr.getBlock() == branch.getBlock() && curr.getValue(DeadwoodBranchBlock.FACING) == facing.getOpposite()) {
                        BlockState prev = genLevel.getBlockState(pos.offset(facing.getOpposite().getNormal()));
                        if (prev.getBlock() == branch.getBlock() && prev.getValue(DeadwoodBranchBlock.FACING) == facing.getOpposite()) {
                            continue;
                        }
                    }

                    double dist = baseLog.distSqr(new Vec3i(pos.getX(), pos.getY(), pos.getZ()));
                    float probability = 0.8f;
                    if (facing == Direction.UP) {
                        probability *= 1.5;

                        // Stop branches from growing from the top of logs
                        if (genLevel.getBlockState(pos).getBlock() == log.getBlock()) {
                            probability = 0;
                        }
                    } else if (facing == Direction.DOWN) {
                        probability = 0.00f;
                    }

                    probability -= 0.05 * dist;
                    probability -= 0.002 * count;
                    probability -= 0.1 * branchLength;
                    probability += 0.25 * balancingCount.get(facing);

                    if (random.nextFloat() < probability) {
                        BlockPos nextPos = pos.offset(facing.getNormal());
                        if (genLevel.isEmptyBlock(nextPos) && genLevel.isAreaLoaded(nextPos, 8)) {
                            DeadwoodBranchBlock branchBlock = (DeadwoodBranchBlock) branch.getBlock();
                            this.setBlock(simulatedRW, nextPos, branchBlock.makeConnections(genLevel, nextPos, facing));
                            placedBranches.add(nextPos);

                            // Add this branch onto the queue to spawn new branches from
                            queue.add(new ImmutablePair<>(nextPos, branchLength + 1));

                            // Keep track of how many branches have been generated
                            count++;

                            // Store the chosen facing for this branch to prevent odd splitting
                            choosenFacings.add(facing);

                            // Store the count for the direction chosen to help with balancing the tree
                            if (facing != Direction.UP && facing != Direction.DOWN) {
                                if (balancingCount.get(facing) > 0) {
                                    balancingCount.put(facing, balancingCount.get(facing) - 1);
                                } else {
                                    balancingCount.put(facing.getOpposite(), balancingCount.get(facing.getOpposite()) + 1);
                                }
                            }
                        }
                    }
                }

                for (BlockPos placedLocation : placedBranches) {
                    if (!simulatedRW.isStateAtPosition(placedLocation, BlockState::isAir)) {
                        DeadwoodBranchBlock branchBlock = (DeadwoodBranchBlock) branch.getBlock();
                        this.setBlock(simulatedRW, placedLocation, branchBlock.makeConnections(genLevel, placedLocation));
                    }
                }

                // Failsafe, if the stack has gotten this large, something has probably gone wrong.
                if (queue.size() > 100) {
                    break;
                }
            }
        }
    }
}