TiedUp-/src/main/java/com/tiedup/remake/blocks/entity/TrappedChestBlockEntity.java

package com.tiedup.remake.blocks.entity;

import com.tiedup.remake.items.clothes.GenericClothes;
import com.tiedup.remake.v2.BodyRegionV2;
import com.tiedup.remake.v2.bondage.BindModeHelper;
import com.tiedup.remake.v2.bondage.CollarHelper;
import com.tiedup.remake.v2.bondage.component.ComponentType;
import com.tiedup.remake.v2.bondage.component.GaggingComponent;
import com.tiedup.remake.v2.bondage.datadriven.DataDrivenBondageItem;
import com.tiedup.remake.v2.bondage.datadriven.DataDrivenItemDefinition;
import com.tiedup.remake.v2.bondage.datadriven.DataDrivenItemRegistry;
import javax.annotation.Nullable;
import net.minecraft.core.BlockPos;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.network.protocol.Packet;
import net.minecraft.network.protocol.game.ClientGamePacketListener;
import net.minecraft.network.protocol.game.ClientboundBlockEntityDataPacket;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.level.block.entity.ChestBlockEntity;
import net.minecraft.world.level.block.state.BlockState;

/**
 * BlockEntity for trapped chest blocks.
 *
 *
 * Extends ChestBlockEntity for proper chest behavior,
 * but also stores bondage items for the trap.
 */
public class TrappedChestBlockEntity
    extends ChestBlockEntity
    implements IBondageItemHolder
{

    // Bondage item storage (separate from chest inventory)
    private ItemStack bind = ItemStack.EMPTY;
    private ItemStack gag = ItemStack.EMPTY;
    private ItemStack blindfold = ItemStack.EMPTY;
    private ItemStack earplugs = ItemStack.EMPTY;
    private ItemStack collar = ItemStack.EMPTY;
    private ItemStack clothes = ItemStack.EMPTY;

    public TrappedChestBlockEntity(BlockPos pos, BlockState state) {
        super(ModBlockEntities.TRAPPED_CHEST.get(), pos, state);
    }

    // BONDAGE ITEM HOLDER IMPLEMENTATION

    @Override
    public ItemStack getBind() {
        return bind;
    }

    @Override
    public void setBind(ItemStack stack) {
        if (stack.isEmpty() || BindModeHelper.isBindItem(stack)) {
            this.bind = stack;
            setChangedAndSync();
        }
    }

    @Override
    public ItemStack getGag() {
        return gag;
    }

    @Override
    public void setGag(ItemStack stack) {
        if (stack.isEmpty()
                || DataDrivenBondageItem.getComponent(stack, ComponentType.GAGGING, GaggingComponent.class) != null) {
            this.gag = stack;
            setChangedAndSync();
        }
    }

    @Override
    public ItemStack getBlindfold() {
        return blindfold;
    }

    @Override
    public void setBlindfold(ItemStack stack) {
        if (stack.isEmpty() || isDataDrivenForRegion(stack, BodyRegionV2.EYES)) {
            this.blindfold = stack;
            setChangedAndSync();
        }
    }

    @Override
    public ItemStack getEarplugs() {
        return earplugs;
    }

    @Override
    public void setEarplugs(ItemStack stack) {
        if (stack.isEmpty() || isDataDrivenForRegion(stack, BodyRegionV2.EARS)) {
            this.earplugs = stack;
            setChangedAndSync();
        }
    }

    @Override
    public ItemStack getCollar() {
        return collar;
    }

    @Override
    public void setCollar(ItemStack stack) {
        if (stack.isEmpty() || CollarHelper.isCollar(stack)) {
            this.collar = stack;
            setChangedAndSync();
        }
    }

    @Override
    public ItemStack getClothes() {
        return clothes;
    }

    @Override
    public void setClothes(ItemStack stack) {
        if (stack.isEmpty() || stack.getItem() instanceof GenericClothes) {
            this.clothes = stack;
            setChangedAndSync();
        }
    }

    @Override
    public boolean isArmed() {
        return (
            !bind.isEmpty() ||
            !gag.isEmpty() ||
            !blindfold.isEmpty() ||
            !earplugs.isEmpty() ||
            !collar.isEmpty() ||
            !clothes.isEmpty()
        );
    }

    @Override
    public void readBondageData(CompoundTag tag) {
        if (tag.contains("bind")) bind = ItemStack.of(tag.getCompound("bind"));
        if (tag.contains("gag")) gag = ItemStack.of(tag.getCompound("gag"));
        if (tag.contains("blindfold")) blindfold = ItemStack.of(
            tag.getCompound("blindfold")
        );
        if (tag.contains("earplugs")) earplugs = ItemStack.of(
            tag.getCompound("earplugs")
        );
        if (tag.contains("collar")) collar = ItemStack.of(
            tag.getCompound("collar")
        );
        if (tag.contains("clothes")) clothes = ItemStack.of(
            tag.getCompound("clothes")
        );
    }

    @Override
    public CompoundTag writeBondageData(CompoundTag tag) {
        if (!bind.isEmpty()) tag.put("bind", bind.save(new CompoundTag()));
        if (!gag.isEmpty()) tag.put("gag", gag.save(new CompoundTag()));
        if (!blindfold.isEmpty()) tag.put(
            "blindfold",
            blindfold.save(new CompoundTag())
        );
        if (!earplugs.isEmpty()) tag.put(
            "earplugs",
            earplugs.save(new CompoundTag())
        );
        if (!collar.isEmpty()) tag.put(
            "collar",
            collar.save(new CompoundTag())
        );
        if (!clothes.isEmpty()) tag.put(
            "clothes",
            clothes.save(new CompoundTag())
        );
        return tag;
    }

    // NBT SERIALIZATION

    @Override
    public void load(CompoundTag tag) {
        super.load(tag);
        readBondageData(tag);
    }

    @Override
    protected void saveAdditional(CompoundTag tag) {
        super.saveAdditional(tag);
        writeBondageData(tag);
    }

    // V2 HELPERS

    /** Check if a stack is a data-driven item occupying the given body region. */
    private static boolean isDataDrivenForRegion(ItemStack stack, BodyRegionV2 region) {
        DataDrivenItemDefinition def = DataDrivenItemRegistry.get(stack);
        return def != null && def.occupiedRegions().contains(region);
    }

    // NETWORK SYNC

    /**
     * Mark dirty and sync to clients.
     * Ensures bondage trap state is visible to all players.
     */
    protected void setChangedAndSync() {
        if (this.level != null) {
            this.setChanged();
            // Notify clients of block update
            this.level.sendBlockUpdated(
                this.worldPosition,
                this.getBlockState(),
                this.getBlockState(),
                3
            );
        }
    }

    @Override
    public CompoundTag getUpdateTag() {
        CompoundTag tag = super.getUpdateTag();
        writeBondageData(tag);
        return tag;
    }

    @Nullable
    @Override
    public Packet<ClientGamePacketListener> getUpdatePacket() {
        return ClientboundBlockEntityDataPacket.create(this);
    }

    @Override
    public void handleUpdateTag(CompoundTag tag) {
        super.handleUpdateTag(tag);
        readBondageData(tag);
    }
}