TravelersSuitcase/src/main/java/io/lampnet/travelerssuitcase/block/entity/SuitcaseBlockEntity.java

package io.lampnet.travelerssuitcase.block.entity;

import io.lampnet.travelerssuitcase.block.PocketPortalBlock;
import net.minecraft.core.BlockPos;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.ListTag;
import net.minecraft.nbt.Tag;
import net.minecraft.network.protocol.game.ClientboundBlockEntityDataPacket;
import net.minecraft.server.level.ServerPlayer;
import net.minecraft.world.level.Level;
import net.minecraft.world.level.block.Block;
import net.minecraft.world.level.block.entity.BlockEntity;
import net.minecraft.world.level.block.state.BlockState;
import org.jetbrains.annotations.Nullable;

import java.util.*;

public class SuitcaseBlockEntity extends BlockEntity {
    private String boundKeystoneName;
    private boolean isLocked = false;
    private boolean dimensionLocked = true;
    private final List<EnteredPlayerData> enteredPlayers = new ArrayList<>();

    public static class EnteredPlayerData {
        public final String uuid;
        public final double x, y, z;
        public final float pitch, yaw;
        public final BlockPos suitcasePos;

        public EnteredPlayerData(String uuid, double x, double y, double z, float pitch, float yaw, BlockPos suitcasePos) {
            this.uuid = uuid; this.x = x; this.y = y; this.z = z;
            this.pitch = pitch; this.yaw = yaw;
            this.suitcasePos = suitcasePos;
        }

        public CompoundTag toNbt() {
            CompoundTag nbt = new CompoundTag();
            nbt.putString("UUID", uuid);
            nbt.putDouble("X", x); nbt.putDouble("Y", y); nbt.putDouble("Z", z);
            nbt.putFloat("Pitch", pitch); nbt.putFloat("Yaw", yaw);
            if (suitcasePos != null) {
                nbt.putInt("SuitcaseX", suitcasePos.getX());
                nbt.putInt("SuitcaseY", suitcasePos.getY());
                nbt.putInt("SuitcaseZ", suitcasePos.getZ());
            }
            return nbt;
        }

        public static EnteredPlayerData fromNbt(CompoundTag nbt) {
            BlockPos sPos = null;
            if (nbt.contains("SuitcaseX") && nbt.contains("SuitcaseY") && nbt.contains("SuitcaseZ")) {
                 sPos = new BlockPos(nbt.getInt("SuitcaseX"), nbt.getInt("SuitcaseY"), nbt.getInt("SuitcaseZ"));
            }
            return new EnteredPlayerData(
                    nbt.getString("UUID"),
                    nbt.getDouble("X"), nbt.getDouble("Y"), nbt.getDouble("Z"),
                    nbt.getFloat("Pitch"), nbt.getFloat("Yaw"),
                    sPos
            );
        }
    }

    public SuitcaseBlockEntity(BlockPos pos, BlockState state) {
        super(ModBlockEntities.SUITCASE_BLOCK_ENTITY.get(), pos, state);
    }

    public boolean canOpenInDimension(Level level) {
        if (!dimensionLocked) {
            return true;
        }
        return level.dimension() == Level.OVERWORLD;
    }

    public boolean isDimensionLocked() {
        return dimensionLocked;
    }

    public void playerEntered(ServerPlayer player) {
        enteredPlayers.removeIf(data -> data.uuid.equals(player.getUUID().toString()));
        EnteredPlayerData data = new EnteredPlayerData(
                player.getUUID().toString(),
                player.getX(), player.getY(), player.getZ(),
                player.getXRot(), player.getYRot(),
                this.worldPosition
        );
        enteredPlayers.add(data);
        if (boundKeystoneName != null) {
            Map<String, BlockPos> suitcases = SUITCASE_REGISTRY.computeIfAbsent(
                    boundKeystoneName, k -> new HashMap<>()
            );
            suitcases.put(player.getUUID().toString(), this.worldPosition);
        }
        PocketPortalBlock.storePlayerPosition(player);
        setChangedAndNotify();
    }

    public EnteredPlayerData getExitPosition(String playerUuid) {
        for (int i = 0; i < enteredPlayers.size(); i++) {
            EnteredPlayerData data = enteredPlayers.get(i);
            if (data.uuid.equals(playerUuid)) {
                EnteredPlayerData exitData = new EnteredPlayerData(
                        data.uuid,
                        this.worldPosition.getX() + 0.5, this.worldPosition.getY() + 1.0, this.worldPosition.getZ() + 0.5,
                        data.pitch, data.yaw,
                        this.worldPosition
                );
                enteredPlayers.remove(i);
                setChangedAndNotify();
                return exitData;
            }
        }
        return null;
    }

    public void bindKeystone(String keystoneName) {
        this.boundKeystoneName = keystoneName;
        setChangedAndNotify();
    }

    public String getBoundKeystoneName() {
        return boundKeystoneName;
    }

    public void setLocked(boolean locked) {
        this.isLocked = locked;
        setChangedAndNotify();
    }

    public boolean isLocked() {
        return isLocked;
    }

    @Override
    public void saveAdditional(CompoundTag nbt) {
        super.saveAdditional(nbt);
        if (boundKeystoneName != null) {
            nbt.putString("BoundKeystone", boundKeystoneName);
        }
        nbt.putBoolean("Locked", isLocked);
        nbt.putBoolean("DimensionLocked", dimensionLocked);

        ListTag playersListTag = new ListTag();
        for (EnteredPlayerData data : enteredPlayers) {
            playersListTag.add(data.toNbt());
        }
        if (!playersListTag.isEmpty()) {
            nbt.put("EnteredPlayers", playersListTag);
        }
    }

    @Override
    public void load(CompoundTag nbt) {
        super.load(nbt);
        if (nbt.contains("BoundKeystone", Tag.TAG_STRING)) {
            boundKeystoneName = nbt.getString("BoundKeystone");
        } else {
            boundKeystoneName = null;
        }
        isLocked = nbt.getBoolean("Locked");
        dimensionLocked = !nbt.contains("DimensionLocked") || nbt.getBoolean("DimensionLocked");
        enteredPlayers.clear();
        if (nbt.contains("EnteredPlayers", Tag.TAG_LIST)) {
            ListTag playersListTag = nbt.getList("EnteredPlayers", Tag.TAG_COMPOUND);
            for (int i = 0; i < playersListTag.size(); i++) {
                enteredPlayers.add(EnteredPlayerData.fromNbt(playersListTag.getCompound(i)));
            }
        }
    }

    @Nullable
    @Override
    public ClientboundBlockEntityDataPacket getUpdatePacket() {
        return ClientboundBlockEntityDataPacket.create(this);
    }

    @Override
    public CompoundTag getUpdateTag() {
        CompoundTag nbt = new CompoundTag();
        this.saveAdditional(nbt);
        return nbt;
    }

    public static final Map<String, Map<String, BlockPos>> SUITCASE_REGISTRY = Collections.synchronizedMap(new HashMap<>());

    public static BlockPos findSuitcasePosition(String keystoneName, String playerUuid) {
        Map<String, BlockPos> suitcases = SUITCASE_REGISTRY.get(keystoneName);
        return (suitcases != null) ? suitcases.get(playerUuid) : null;
    }

    public static void removeSuitcaseEntry(String keystoneName, String playerUuid) {
        Map<String, BlockPos> suitcases = SUITCASE_REGISTRY.get(keystoneName);
        if (suitcases != null) {
            suitcases.remove(playerUuid);
            if (suitcases.isEmpty()) {
                SUITCASE_REGISTRY.remove(keystoneName);
            }
        }
    }

    public List<EnteredPlayerData> getEnteredPlayers() {
        return Collections.unmodifiableList(enteredPlayers);
    }

    public static void initializeSuitcaseRegistry(Map<String, Map<String, BlockPos>> savedRegistry) {
        SUITCASE_REGISTRY.clear();
        if (savedRegistry != null) {
            savedRegistry.forEach((key, value) -> {
                Map<String, BlockPos> players = SUITCASE_REGISTRY.computeIfAbsent(key, k -> Collections.synchronizedMap(new HashMap<>()));
                players.putAll(value);
            });
        }
    }

    public static void saveSuitcaseRegistryTo(Map<String, Map<String, BlockPos>> destination) {
        destination.clear();
        SUITCASE_REGISTRY.forEach((key, value) -> {
             Map<String, BlockPos> players = destination.computeIfAbsent(key, k -> new HashMap<>());
             players.putAll(value);
        });
    }

    public void updatePlayerSuitcasePosition(String playerUuid, BlockPos newPos) {
        boolean updated = false;
        for (int i = 0; i < enteredPlayers.size(); i++) {
            EnteredPlayerData data = enteredPlayers.get(i);
            if (data.uuid.equals(playerUuid)) {
                enteredPlayers.set(i, new EnteredPlayerData(data.uuid, data.x, data.y, data.z, data.pitch, data.yaw, newPos));
                updated = true;
                break;
            }
        }
        if (updated && boundKeystoneName != null) {
            Map<String, BlockPos> suitcases = SUITCASE_REGISTRY.computeIfAbsent(boundKeystoneName, k -> new HashMap<>());
            suitcases.put(playerUuid, newPos);
        }
        if(updated) setChangedAndNotify();
    }

    private void setChangedAndNotify() {
        setChanged();
        if (level != null && !level.isClientSide()) {
            level.sendBlockUpdated(worldPosition, getBlockState(), getBlockState(), Block.UPDATE_ALL);
        }
    }
}