From 5f066b007a5058c6e65b9c4d556fd805f2551fb3 Mon Sep 17 00:00:00 2001
From: DBotThePony <dbotthepony@yandex.ru>
Date: Sun, 8 Aug 2021 15:44:18 +0700
Subject: [PATCH] more capability tests

---
 .../mc/otm/block/BlockAndroidStation.java     | 59 ++++++++++++++++++-
 .../mc/otm/capability/AndroidCapability.java  |  5 ++
 .../mc/otm/capability/IAndroidCapability.java | 11 ++++
 3 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/src/main/java/ru/dbotthepony/mc/otm/block/BlockAndroidStation.java b/src/main/java/ru/dbotthepony/mc/otm/block/BlockAndroidStation.java
index f4383402b..fea3ed8c1 100644
--- a/src/main/java/ru/dbotthepony/mc/otm/block/BlockAndroidStation.java
+++ b/src/main/java/ru/dbotthepony/mc/otm/block/BlockAndroidStation.java
@@ -5,6 +5,7 @@ import net.minecraft.core.Direction;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.server.level.ServerPlayer;
+import net.minecraft.world.entity.LivingEntity;
 import net.minecraft.world.level.BlockGetter;
 import net.minecraft.world.level.Level;
 import net.minecraft.world.level.block.Block;
@@ -22,8 +23,13 @@ import net.minecraft.world.phys.shapes.Shapes;
 import net.minecraft.world.phys.shapes.VoxelShape;
 import net.minecraftforge.common.capabilities.Capability;
 import net.minecraftforge.common.util.LazyOptional;
+import net.minecraftforge.energy.CapabilityEnergy;
+import net.minecraftforge.energy.EnergyStorage;
+import net.minecraftforge.energy.IEnergyStorage;
 import org.apache.logging.log4j.core.jmx.Server;
 import ru.dbotthepony.mc.otm.OverdriveThatMatters;
+import ru.dbotthepony.mc.otm.capability.AndroidCapability;
+import ru.dbotthepony.mc.otm.capability.IAndroidCapability;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -56,27 +62,76 @@ public class BlockAndroidStation extends Block implements EntityBlock {
 	}
 
 	public static class BlockAndroidStationEntity extends BlockEntity {
+		private static class AndroidStationEnergyStorage extends EnergyStorage {
+			public AndroidStationEnergyStorage(int capacity, int maxReceive, int maxExtract) {
+				super(capacity, maxReceive, maxExtract);
+			}
+
+			void setEnergy(int energy) {
+				this.energy = energy;
+			}
+		}
+
+		private final AndroidStationEnergyStorage energy = new AndroidStationEnergyStorage(64000, 1024, 1024);
+
 		public BlockAndroidStationEntity(BlockPos p_155229_, BlockState p_155230_) {
 			super(OverdriveThatMatters.ANDROID_STATION_FACTORY, p_155229_, p_155230_);
 		}
 
 		@Nonnull
 		public <T> LazyOptional<T> getCapability(@Nonnull Capability<T> cap, @Nullable Direction side) {
+			if (cap == CapabilityEnergy.ENERGY) {
+				return LazyOptional.of(() -> energy).cast();
+			}
+
 			return super.getCapability(cap, side);
 		}
 
+		@Override
+		public CompoundTag save(CompoundTag nbt) {
+			nbt.putInt("energy_stored", energy.getEnergyStored());
+			return super.save(nbt);
+		}
+
 		public void load(CompoundTag nbt) {
 			super.load(nbt);
 
+			if (nbt.contains("energy_stored"))
+				energy.setEnergy(nbt.getInt("energy_stored"));
 		}
 
 		void tick() {
 			if (!hasLevel())
 				return;
 
+			if (energy.getEnergyStored() <= 0)
+				return;
+
 			BlockPos pos = getBlockPos();
-			List<ServerPlayer> players = getLevel().getEntitiesOfClass(ServerPlayer.class, new AABB(pos.getX(), pos.getY(), pos.getZ(), pos.getX() + 1, pos.getY() + 2, pos.getZ() + 1));
-			OverdriveThatMatters.LOGGER.info("Players standing above {}", players);
+			List<LivingEntity> entities = getLevel().getEntitiesOfClass(LivingEntity.class, new AABB(pos.getX(), pos.getY(), pos.getZ(), pos.getX() + 1, pos.getY() + 2, pos.getZ() + 1));
+
+			for (LivingEntity ent : entities) {
+				LazyOptional<IAndroidCapability> resolver = ent.getCapability(AndroidCapability.CAPABILITY);
+
+				if (resolver.isPresent()) {
+					IAndroidCapability capability = resolver.resolve().get();
+
+					int missing = capability.getMissingPowerInt();
+
+					if (missing > 0) {
+						int extract = energy.extractEnergy(missing, true);
+
+						if (extract > 0) {
+							int received = (int) capability.receiveEnergy(extract, true);
+
+							if (received > 0) {
+								energy.extractEnergy(extract, false);
+								capability.receiveEnergy(extract, false);
+							}
+						}
+					}
+				}
+			}
 		}
 
 		public static <T extends BlockEntity> void tick(Level level, BlockPos blockPos, BlockState blockState, T t) {
diff --git a/src/main/java/ru/dbotthepony/mc/otm/capability/AndroidCapability.java b/src/main/java/ru/dbotthepony/mc/otm/capability/AndroidCapability.java
index 5ce5b0500..6c778fc90 100644
--- a/src/main/java/ru/dbotthepony/mc/otm/capability/AndroidCapability.java
+++ b/src/main/java/ru/dbotthepony/mc/otm/capability/AndroidCapability.java
@@ -52,6 +52,11 @@ public class AndroidCapability implements ICapabilityProvider, IAndroidCapabilit
 		}
 	}
 
+	@Override
+	public long getMissingPower() {
+		return getMaxBatteryLevel() - getBatteryLevel();
+	}
+
 	@Override
 	public void tick() {
 		if (ent instanceof ServerPlayer) {
diff --git a/src/main/java/ru/dbotthepony/mc/otm/capability/IAndroidCapability.java b/src/main/java/ru/dbotthepony/mc/otm/capability/IAndroidCapability.java
index 11132e67f..f8ff74ee3 100644
--- a/src/main/java/ru/dbotthepony/mc/otm/capability/IAndroidCapability.java
+++ b/src/main/java/ru/dbotthepony/mc/otm/capability/IAndroidCapability.java
@@ -9,6 +9,17 @@ public interface IAndroidCapability {
 
 	long getBatteryLevel();
 	long getMaxBatteryLevel();
+	long getMissingPower();
 	long receiveEnergy(long howMuch, boolean simulate);
 	long extractEnergy(long howMuch, boolean simulate);
+
+	default int getMissingPowerInt() {
+		long power = getMissingPower();
+
+		if (power > Integer.MAX_VALUE) {
+			return Integer.MAX_VALUE;
+		}
+
+		return (int) power;
+	}
 }