EnergonRelics/src/main/java/com/thebrokenrail/energonrelics/block/entity/LightningRodBaseBlockEntity.java

package com.thebrokenrail.energonrelics.block.entity;

import com.thebrokenrail.energonrelics.api.block.entity.core.EnergyProviderBlockEntity;
import com.thebrokenrail.energonrelics.api.block.entity.helper.EnergyGenerator;
import com.thebrokenrail.energonrelics.api.energy.Action;
import com.thebrokenrail.energonrelics.config.HardcodedConfig;
import com.thebrokenrail.energonrelics.api.energy.tick.EnergyTickable;
import net.minecraft.block.BlockState;
import net.minecraft.block.entity.BlockEntityType;
import net.minecraft.entity.EntityType;
import net.minecraft.entity.LightningEntity;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.world.LightType;

import java.util.List;
import java.util.Objects;

public class LightningRodBaseBlockEntity extends EnergyProviderBlockEntity implements EnergyGenerator {
    public LightningRodBaseBlockEntity(BlockEntityType<?> type) {
        super(type);
    }

    private int getLight() {
        if (hasWorld() && Objects.requireNonNull(getWorld()).getDimension().hasSkyLight() && getWorld().isThundering()) {
            assert getWorld() != null;
            return getWorld().getLightLevel(LightType.SKY, getPos());
        } else {
            return 0;
        }
    }

    private int cooldown = 0;
    private long energy = 0;

    @Override
    public long getDisplayEnergy() {
        return energy;
    }

    @Override
    public long getEnergy() {
        return energy;
    }

    @Override
    public void setEnergy(long energy) {
        this.energy = energy;
    }

    @Override
    public List<EnergyTickable> startTick() {
        assert getWorld() != null;

        if (cooldown <= 0) {
            boolean dirty = false;

            if (energy != 0) {
                energy = 0;
                dirty = true;
            }

            if (getLight() > 0) {
                if (getWorld().random.nextDouble() <= HardcodedConfig.LIGHTNING_ROD_CHANCE) {
                    LightningEntity entity = EntityType.LIGHTNING_BOLT.create(getWorld());
                    assert entity != null;
                    entity.updatePosition(getPos().getX() + 0.5d, getPos().getY() + 1d + (getWorld().random.nextDouble() * 1.5d), getPos().getZ() + 0.5d);
                    entity.setCosmetic(true);
                    Objects.requireNonNull(getWorld()).spawnEntity(entity);

                    cooldown = HardcodedConfig.LIGHTNING_ROD_COOLDOWN;
                    energy = HardcodedConfig.LIGHTNING_ROD_ENERGY_OUTPUT;

                    dirty = true;
                }
            }

            if (dirty) {
                markDirty();
            }
        } else {
            energy = 0;

            cooldown--;
            markDirty();
        }

        return super.startTick();
    }

    @Override
    public CompoundTag toTag(CompoundTag tag) {
        super.toTag(tag);
        tag.putInt("Cooldown", cooldown);
        return tag;
    }

    @Override
    public void fromTag(BlockState state, CompoundTag tag) {
        super.fromTag(state, tag);
        cooldown = tag.getInt("Cooldown");
    }

    @Override
    protected void handlePropagatedAction(Action.PropagatedAction action) {
        super.handlePropagatedAction(action);
        handlePropagatedActionWithGenerator(action);
    }

    @Override
    public boolean isEnergyProvider() {
        return true;
    }
}