TiedUp-/src/main/java/com/tiedup/remake/events/lifecycle/PlayerStateEventHandler.java

package com.tiedup.remake.events.lifecycle;

import com.tiedup.remake.cells.CampOwnership;
import com.tiedup.remake.cells.CellRegistryV2;
import com.tiedup.remake.core.TiedUpMod;
import com.tiedup.remake.events.restriction.BondageItemRestrictionHandler;
import com.tiedup.remake.labor.LaborTask;
import com.tiedup.remake.network.ModNetwork;
import com.tiedup.remake.network.labor.PacketSyncLaborProgress;
import com.tiedup.remake.network.sync.PacketSyncBindState;
import com.tiedup.remake.prison.LaborRecord;
import com.tiedup.remake.prison.PrisonerManager;
import com.tiedup.remake.prison.PrisonerRecord;
import com.tiedup.remake.prison.PrisonerState;
import com.tiedup.remake.state.IPlayerLeashAccess;
import com.tiedup.remake.state.PlayerBindState;
import com.tiedup.remake.v2.BodyRegionV2;
import com.tiedup.remake.v2.bondage.capability.V2BondageEquipmentProvider;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import net.minecraft.core.BlockPos;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.resources.ResourceKey;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.server.level.ServerPlayer;
import net.minecraft.world.entity.decoration.LeashFenceKnotEntity;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.level.Level;
import net.minecraftforge.event.entity.living.LivingDeathEvent;
import net.minecraftforge.event.entity.player.PlayerEvent;
import net.minecraftforge.eventbus.api.SubscribeEvent;
import net.minecraftforge.fml.common.Mod;

/**
 * Event handler for PlayerBindState lifecycle management.
 * Handles player connection, disconnection, death, and respawn.
 */
@Mod.EventBusSubscriber(
    modid = TiedUpMod.MOD_ID,
    bus = Mod.EventBusSubscriber.Bus.FORGE
)
public class PlayerStateEventHandler {

    /** UUIDs of players who died while imprisoned — message sent on respawn */
    private static final Set<UUID> pendingDeathEscapeMessage =
        ConcurrentHashMap.newKeySet();

    /**
     * Called when a player logs in to the server.
     * Initialize or reset their PlayerBindState, restore pole leash, and sync to client.
     */
    @SubscribeEvent
    public static void onPlayerLoggedIn(PlayerEvent.PlayerLoggedInEvent event) {
        if (!(event.getEntity() instanceof ServerPlayer player)) {
            return;
        }

        PlayerBindState state = PlayerBindState.getInstance(player);
        if (state != null) {
            state.resetNewConnection(player);
        }

        // Restore pole leash if player was leashed to one when they disconnected
        restorePoleLeashIfNeeded(player);

        // MEDIUM FIX: Restore captor if player was captive when they disconnected
        restoreCaptorIfNeeded(player);

        // Sync V2 equipment + bind state to the logging-in player
        com.tiedup.remake.v2.bondage.capability.V2EquipmentHelper.sync(player);

        PacketSyncBindState statePacket = PacketSyncBindState.fromPlayer(
            player
        );
        if (statePacket != null) {
            ModNetwork.sendToPlayer(statePacket, player);
        }

        // Sync Labor HUD on login
        syncLaborState(player);
    }

    /**
     * Restore leash to pole if player was leashed when they disconnected.
     */
    private static void restorePoleLeashIfNeeded(ServerPlayer player) {
        player
            .getCapability(V2BondageEquipmentProvider.V2_BONDAGE_EQUIPMENT)
            .ifPresent(cap -> {
                if (!cap.wasLeashedToPole()) {
                    return;
                }

                BlockPos polePos = cap.getSavedPolePosition();
                ResourceKey<Level> poleDimension = cap.getSavedPoleDimension();

                if (polePos == null || poleDimension == null) {
                    cap.clearSavedPoleLeash();
                    return;
                }

                // Check if player is in the same dimension
                if (!player.level().dimension().equals(poleDimension)) {
                    TiedUpMod.LOGGER.info(
                        "[Lifecycle] {} reconnected but is in different dimension. Clearing saved pole leash.",
                        player.getName().getString()
                    );
                    cap.clearSavedPoleLeash();
                    return;
                }

                ServerLevel level = player.serverLevel();

                // Try to find or create the LeashFenceKnotEntity at that position
                LeashFenceKnotEntity fenceKnot =
                    LeashFenceKnotEntity.getOrCreateKnot(level, polePos);

                if (fenceKnot == null) {
                    TiedUpMod.LOGGER.info(
                        "[Lifecycle] {} reconnected but pole at {} no longer exists.",
                        player.getName().getString(),
                        polePos
                    );
                    cap.clearSavedPoleLeash();
                    return;
                }

                // Attach leash to the pole
                if (player instanceof IPlayerLeashAccess access) {
                    access.tiedup$attachLeash(fenceKnot);
                    TiedUpMod.LOGGER.info(
                        "[Lifecycle] {} reconnected. Restored leash to pole at {}.",
                        player.getName().getString(),
                        polePos
                    );
                }

                // Clear saved data (no longer needed)
                cap.clearSavedPoleLeash();
            });
    }

    /**
     * MEDIUM FIX: Restore captor if player was captive when they disconnected.
     */
    private static void restoreCaptorIfNeeded(ServerPlayer player) {
        player
            .getCapability(V2BondageEquipmentProvider.V2_BONDAGE_EQUIPMENT)
            .ifPresent(cap -> {
                if (!cap.hasSavedCaptor()) {
                    return;
                }

                java.util.UUID captorUUID = cap.getSavedCaptorUUID();
                if (captorUUID == null) {
                    cap.clearSavedCaptor();
                    return;
                }

                ServerLevel level = player.serverLevel();

                // Try to find the captor entity
                net.minecraft.world.entity.Entity captorEntity =
                    level.getEntity(captorUUID);

                if (captorEntity == null) {
                    TiedUpMod.LOGGER.info(
                        "[Lifecycle] {} reconnected but captor {} no longer exists. Clearing captivity.",
                        player.getName().getString(),
                        captorUUID.toString().substring(0, 8)
                    );
                    cap.clearSavedCaptor();
                    return;
                }

                // Check if captor is still a valid ICaptor
                if (
                    !(captorEntity instanceof
                            net.minecraft.world.entity.LivingEntity livingCaptor)
                ) {
                    TiedUpMod.LOGGER.info(
                        "[Lifecycle] {} reconnected but captor is not a LivingEntity. Clearing captivity.",
                        player.getName().getString()
                    );
                    cap.clearSavedCaptor();
                    return;
                }

                // Get the ICaptor interface
                com.tiedup.remake.state.ICaptor captor = null;
                if (
                    livingCaptor instanceof
                        com.tiedup.remake.state.ICaptor kidnapper
                ) {
                    captor = kidnapper;
                } else {
                    TiedUpMod.LOGGER.info(
                        "[Lifecycle] {} reconnected but captor {} does not implement ICaptor. Clearing captivity.",
                        player.getName().getString(),
                        captorUUID.toString().substring(0, 8)
                    );
                    cap.clearSavedCaptor();
                    return;
                }

                // Check if player is still enslavable (still tied up)
                PlayerBindState state = PlayerBindState.getInstance(player);
                if (state == null || !state.isTiedUp()) {
                    TiedUpMod.LOGGER.info(
                        "[Lifecycle] {} reconnected but is no longer tied up. Clearing saved captor.",
                        player.getName().getString()
                    );
                    cap.clearSavedCaptor();
                    return;
                }

                // Restore captivity using the PlayerCaptivity component
                boolean success = state.getCapturedBy(captor);

                if (success) {
                    TiedUpMod.LOGGER.info(
                        "[Lifecycle] {} reconnected. Restored captivity to {}.",
                        player.getName().getString(),
                        captorEntity.getName().getString()
                    );
                } else {
                    TiedUpMod.LOGGER.warn(
                        "[Lifecycle] {} reconnected but failed to restore captivity to {}.",
                        player.getName().getString(),
                        captorEntity.getName().getString()
                    );
                }

                // Clear saved data (no longer needed)
                cap.clearSavedCaptor();
            });
    }

    /**
     * Called when a player logs out of the server.
     * Clean up their PlayerBindState instance.
     */
    @SubscribeEvent
    public static void onPlayerLoggedOut(
        PlayerEvent.PlayerLoggedOutEvent event
    ) {
        if (!(event.getEntity() instanceof ServerPlayer player)) {
            return;
        }

        // NOTE: PlayerBindState.removeInstance() is NOT called here.
        // The canonical removal point is PlayerLifecycleHandler.onPlayerLoggedOut (EventPriority.HIGH).
        // Calling it here at NORMAL priority would be a redundant double-removal.

        // Clean up message cooldowns to prevent memory leak
        BondageItemRestrictionHandler.clearCooldowns(player.getUUID());

        // Clean up pending death escape message flag (player disconnected between death and respawn)
        pendingDeathEscapeMessage.remove(player.getUUID());
    }

    /**
     * Called when a player dies.
     * Handle bondage state cleanup based on game rules.
     */
    @SubscribeEvent
    public static void onPlayerDeath(LivingDeathEvent event) {
        if (!(event.getEntity() instanceof ServerPlayer player)) {
            return;
        }

        PlayerBindState state = PlayerBindState.getInstance(player);
        if (state != null) {
            state.onDeathKidnapped(player.level());
        }

        // Clean up captivity state - prisoner can't continue task after death
        PrisonerManager manager = PrisonerManager.get(player.serverLevel());
        PrisonerRecord record = manager.getRecord(player.getUUID());
        LaborRecord laborRecord = manager.getLaborRecord(player.getUUID());

        if (record.isImprisoned() || laborRecord.hasTask()) {
            // SECURITY: Remove labor tools before death to prevent item duplication
            // (Tools would otherwise drop on death)
            removeLaborToolsOnDeath(player);

            // Transition to FREE state via escape (clears all data)
            long currentTime = player.serverLevel().getGameTime();
            // Use centralized escape service for complete cleanup
            com.tiedup.remake.prison.service.PrisonerService.get().escape(
                player.serverLevel(),
                player.getUUID(),
                "death"
            );

            // Flag for respawn message
            pendingDeathEscapeMessage.add(player.getUUID());

            TiedUpMod.LOGGER.debug(
                "[PlayerStateEventHandler] Transitioned {} to FREE state on death",
                player.getName().getString()
            );
        } else if (record.isProtected(player.serverLevel().getGameTime())) {
            // Revoke grace period on death - player can be targeted again after respawn
            // Death is considered "payment" for any remaining grace
            record.setProtectionExpiry(0);
        }

        // Release prisoner from any cells to prevent "ghost prisoner" blocking cells
        // Without this, the cell still counts this player as prisoner, making it "full"
        // and preventing new kidnappers from using it
        CellRegistryV2 cellRegistry = CellRegistryV2.get(player.server);
        if (cellRegistry != null) {
            int released = cellRegistry.releasePrisonerFromAllCells(
                player.getUUID()
            );
            if (released > 0) {
                TiedUpMod.LOGGER.debug(
                    "[PlayerStateEventHandler] Released {} from {} cell(s) on death",
                    player.getName().getString(),
                    released
                );
            }
        }
    }

    /**
     * Called when a player respawns after death.
     * Handled by PlayerEvent.Clone in CapabilityEventHandler (capability copying)
     * and PlayerLoggedIn (state reset).
     *
     * The respawn process:
     * 1. PlayerEvent.Clone -> Copy capability data from old player to new player
     * 2. PlayerLoggedIn -> Reset PlayerBindState with new player entity
     * 3. Sync to client
     */
    @SubscribeEvent
    public static void onPlayerRespawn(PlayerEvent.PlayerRespawnEvent event) {
        if (!(event.getEntity() instanceof ServerPlayer player)) {
            return;
        }

        // Get or create state for the respawned player
        PlayerBindState state = PlayerBindState.getInstance(player);
        if (state != null) {
            state.resetNewConnection(player);
        }

        // Sync V2 equipment + bind state to the respawned player
        com.tiedup.remake.v2.bondage.capability.V2EquipmentHelper.sync(player);

        PacketSyncBindState statePacket = PacketSyncBindState.fromPlayer(
            player
        );
        if (statePacket != null) {
            ModNetwork.sendToPlayer(statePacket, player);
        }

        // Sync Labor HUD on respawn (restores task OR clears it if executed)
        syncLaborState(player);

        // Send death escape message if applicable
        if (pendingDeathEscapeMessage.remove(player.getUUID())) {
            player.sendSystemMessage(
                net.minecraft.network.chat.Component.translatable(
                    "msg.tiedup.event.death_escape"
                ).withStyle(net.minecraft.ChatFormatting.YELLOW)
            );
        }
    }

    /**
     * Helper to sync labor state to client.
     * Sends current progress if active, or a clear packet if not.
     */
    private static void syncLaborState(ServerPlayer player) {
        PrisonerManager manager = PrisonerManager.get(player.serverLevel());
        PrisonerRecord record = manager.getRecord(player.getUUID());
        LaborRecord laborRecord = manager.getLaborRecord(player.getUUID());

        if (record.isImprisoned() && laborRecord.hasTask()) {
            LaborTask task = laborRecord.getTask();
            ModNetwork.sendToPlayer(
                new PacketSyncLaborProgress(
                    task.getDescription(),
                    task.getProgress(),
                    task.getQuota(),
                    task.getValue()
                ),
                player
            );
            TiedUpMod.LOGGER.debug(
                "[PlayerStateEventHandler] Synced labor HUD for {}: {}/{} {}",
                player.getName().getString(),
                task.getProgress(),
                task.getQuota(),
                task.getDescription()
            );
        } else {
            // No active captivity state - ensure HUD is cleared (important after execution/death)
            ModNetwork.sendToPlayer(new PacketSyncLaborProgress(), player);
            TiedUpMod.LOGGER.debug(
                "[PlayerStateEventHandler] Cleared labor HUD for {}",
                player.getName().getString()
            );
        }
    }

    /**
     * SECURITY: Remove all labor tools from player inventory before death.
     * Prevents tools from dropping and being recoverable.
     */
    private static void removeLaborToolsOnDeath(ServerPlayer player) {
        var inventory = player.getInventory();
        int removedCount = 0;

        for (int i = 0; i < inventory.getContainerSize(); i++) {
            ItemStack stack = inventory.getItem(i);
            if (!stack.isEmpty() && stack.hasTag()) {
                CompoundTag tag = stack.getTag();
                if (tag != null && tag.getBoolean("LaborTool")) {
                    inventory.setItem(i, ItemStack.EMPTY);
                    removedCount++;
                }
            }
        }

        if (removedCount > 0) {
            TiedUpMod.LOGGER.debug(
                "[PlayerStateEventHandler] Removed {} labor tools from {} before death",
                removedCount,
                player.getName().getString()
            );
        }
    }
}