TiedUp-/src/main/java/com/tiedup/remake/prison/service/PrisonerService.java

package com.tiedup.remake.prison.service;

import com.tiedup.remake.cells.CellDataV2;
import com.tiedup.remake.cells.CellRegistryV2;
import com.tiedup.remake.core.TiedUpMod;
import com.tiedup.remake.entities.AbstractTiedUpNpc;
import com.tiedup.remake.entities.EntityDamsel;
import com.tiedup.remake.personality.PersonalityState;
import com.tiedup.remake.prison.PrisonerManager;
import com.tiedup.remake.prison.PrisonerRecord;
import com.tiedup.remake.prison.PrisonerState;
import com.tiedup.remake.state.IBondageState;
import com.tiedup.remake.state.ICaptor;
import com.tiedup.remake.util.KidnappedHelper;
import java.util.*;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.entity.player.Player;
import org.jetbrains.annotations.Nullable;

/**
 * Prisoner lifecycle service.
 *
 * Manages prisoner state transitions atomically:
 * - PrisonerManager (SavedData state)
 * - CellRegistryV2 (cell assignment)
 * - Leash/captor refs (IBondageState / ICaptor)
 *
 * Handles: capture, imprison, extractFromCell, returnToCell, transferCaptive.
 * Escape detection and release are handled by {@link EscapeMonitorService}.
 *
 * AI goals manage: Navigation, Teleport, Equipment, Dialogue.
 * PrisonerService manages: State transitions + leash coordination.
 */
public class PrisonerService {

    private static final PrisonerService INSTANCE = new PrisonerService();

    public static PrisonerService get() {
        return INSTANCE;
    }

    private PrisonerService() {}

    // ==================== CAPTURE ====================

    /**
     * Capture a free entity.
     * FREE → CAPTURED + leash + captor refs.
     *
     * @param level Server level
     * @param captive The entity being captured (must have IBondageState state)
     * @param captor The capturing entity
     * @return true if capture succeeded
     */
    public boolean capture(
        ServerLevel level,
        LivingEntity captive,
        ICaptor captor
    ) {
        UUID captiveId = captive.getUUID();
        PrisonerManager manager = PrisonerManager.get(level);
        long gameTime = level.getGameTime();

        // 1. PrisonerManager state transition: FREE → CAPTURED
        if (
            !manager.capture(captiveId, captor.getEntity().getUUID(), gameTime)
        ) {
            TiedUpMod.LOGGER.warn(
                "[PrisonerService] capture() failed PrisonerManager transition for {}",
                captiveId.toString().substring(0, 8)
            );
            return false;
        }

        // 2. Leash attach
        IBondageState kidnappedState = KidnappedHelper.getKidnappedState(
            captive
        );
        if (kidnappedState == null) {
            // Rollback
            manager.escape(
                captiveId,
                gameTime,
                "capture rollback: no kidnapped state"
            );
            return false;
        }

        boolean leashed;
        if (captive instanceof AbstractTiedUpNpc npc) {
            // NPC: use forceCapturedBy to bypass canCapture() PrisonerManager check
            // (PrisonerManager state is already CAPTURED at this point)
            leashed = npc.forceCapturedBy(captor);
        } else {
            // Player: getCapturedBy checks isEnslavable() first (OR logic),
            // so canCapture() PrisonerManager check is bypassed if tied up.
            // Generic IBondageState: same path.
            leashed = kidnappedState.getCapturedBy(captor);
        }

        if (!leashed) {
            // Rollback PrisonerManager
            manager.escape(
                captiveId,
                gameTime,
                "capture rollback: leash failed"
            );
            TiedUpMod.LOGGER.warn(
                "[PrisonerService] capture() leash failed for {}, rolled back",
                captiveId.toString().substring(0, 8)
            );
            return false;
        }

        TiedUpMod.LOGGER.info(
            "[PrisonerService] Captured {} by {}",
            captive.getName().getString(),
            captor.getEntity().getName().getString()
        );
        return true;
    }

    // ==================== IMPRISONMENT ====================

    /**
     * Imprison a captured entity in a cell.
     * CAPTURED → IMPRISONED + unleash + CellRegistryV2.assign.
     *
     * The caller handles: teleport, confiscation, bind, dialogue.
     *
     * @param level Server level
     * @param captive The captured entity
     * @param captor The current captor
     * @param campId Camp UUID
     * @param cell The target cell
     * @return true if imprisonment succeeded
     */
    public boolean imprison(
        ServerLevel level,
        LivingEntity captive,
        ICaptor captor,
        UUID campId,
        CellDataV2 cell
    ) {
        if (cell == null) {
            TiedUpMod.LOGGER.warn(
                "[PrisonerService] imprison() called with null cell"
            );
            return false;
        }

        UUID captiveId = captive.getUUID();
        PrisonerManager manager = PrisonerManager.get(level);
        CellRegistryV2 cellRegistry = CellRegistryV2.get(level);
        long gameTime = level.getGameTime();

        // 1. Assign to cell registry
        boolean assigned = cellRegistry.assignPrisoner(cell.getId(), captiveId);
        if (!assigned) {
            TiedUpMod.LOGGER.warn(
                "[PrisonerService] imprison() cell assignment failed for {}",
                captiveId.toString().substring(0, 8)
            );
            return false;
        }

        // Defense: if prisoner is FREE, capture first
        boolean defenseCaptured = false;
        PrisonerRecord record = manager.getRecord(captiveId);
        if (record.getState() == PrisonerState.FREE) {
            TiedUpMod.LOGGER.info(
                "[PrisonerService] {} is FREE, capturing before imprison",
                captive.getName().getString()
            );
            manager.capture(captiveId, captor.getEntity().getUUID(), gameTime);
            defenseCaptured = true;
        }

        // 2. PrisonerManager state transition: CAPTURED → IMPRISONED
        boolean imprisoned = manager.imprison(
            captiveId,
            campId,
            cell.getId(),
            gameTime
        );
        if (!imprisoned) {
            // Rollback cell assignment
            cellRegistry.releasePrisoner(
                cell.getId(),
                captiveId,
                level.getServer()
            );
            // Rollback defense auto-capture if we triggered it
            if (defenseCaptured) {
                manager.escape(captiveId, gameTime, "imprison_rollback");
            }
            TiedUpMod.LOGGER.warn(
                "[PrisonerService] imprison() PrisonerManager transition failed for {} (state={})",
                captiveId.toString().substring(0, 8),
                manager.getState(captiveId)
            );
            return false;
        }

        // 3. Unleash: remove captor refs and drop leash
        IBondageState kidnappedState = KidnappedHelper.getKidnappedState(
            captive
        );
        if (kidnappedState != null) {
            kidnappedState.free(false); // Drop leash without dialogue/items
            captor.removeCaptive(kidnappedState, false);
        }

        // 4. Sync PersonalityState cellId for damsels
        if (captive instanceof EntityDamsel damsel) {
            PersonalityState pState = damsel.getPersonalityState();
            if (pState != null) {
                pState.assignCell(cell.getId(), cell, damsel.getUUID());
            }
        }

        TiedUpMod.LOGGER.info(
            "[PrisonerService] Imprisoned {} in cell {} (camp={})",
            captive.getName().getString(),
            cell.getId().toString().substring(0, 8),
            campId != null ? campId.toString().substring(0, 8) : "null"
        );
        return true;
    }

    // ==================== EXTRACTION ====================

    /**
     * Extract a prisoner from their cell.
     * IMPRISONED → WORKING + CellRegistryV2.release + leash to new captor.
     *
     * For: dogwalk, maid fetch, labor extract.
     * Note: If the entity is not tracked by PrisonerManager (e.g., damsels),
     * the PrisonerManager step is skipped — only CellRegistryV2 + leash are handled.
     *
     * @param level Server level
     * @param captive The imprisoned entity
     * @param newCaptor The entity extracting (kidnapper or maid)
     * @param cell The cell to extract from
     * @return true if extraction succeeded
     */
    public boolean extractFromCell(
        ServerLevel level,
        LivingEntity captive,
        ICaptor newCaptor,
        CellDataV2 cell
    ) {
        UUID captiveId = captive.getUUID();
        PrisonerManager manager = PrisonerManager.get(level);
        CellRegistryV2 cellRegistry = CellRegistryV2.get(level);
        long gameTime = level.getGameTime();

        // 1. PrisonerManager: IMPRISONED → WORKING (players only — damsels aren't tracked)
        boolean isPlayer = captive instanceof Player;
        if (isPlayer) {
            if (!manager.extract(captiveId, gameTime)) {
                TiedUpMod.LOGGER.warn(
                    "[PrisonerService] extractFromCell() PrisonerManager transition failed for {}",
                    captiveId.toString().substring(0, 8)
                );
                return false;
            }
        }

        // 2. Release from cell registry
        cellRegistry.releasePrisoner(
            cell.getId(),
            captiveId,
            level.getServer()
        );

        // 3. Leash to new captor — forceCapturedBy for damsels
        IBondageState kidnappedState = KidnappedHelper.getKidnappedState(
            captive
        );
        if (kidnappedState == null) {
            // Rollback
            if (isPlayer) {
                manager.returnToCell(captiveId, cell.getId(), gameTime);
            }
            cellRegistry.assignPrisoner(cell.getId(), captiveId);
            TiedUpMod.LOGGER.warn(
                "[PrisonerService] extractFromCell() no kidnapped state for {}, rolled back",
                captiveId.toString().substring(0, 8)
            );
            return false;
        }

        boolean leashed;
        if (captive instanceof AbstractTiedUpNpc npc) {
            leashed = npc.forceCapturedBy(newCaptor);
        } else {
            leashed = kidnappedState.getCapturedBy(newCaptor);
        }

        if (!leashed) {
            // Rollback: return to IMPRISONED + reassign cell
            if (isPlayer) {
                manager.returnToCell(captiveId, cell.getId(), gameTime);
            }
            cellRegistry.assignPrisoner(cell.getId(), captiveId);
            TiedUpMod.LOGGER.warn(
                "[PrisonerService] extractFromCell() leash failed for {}, rolled back",
                captiveId.toString().substring(0, 8)
            );
            return false;
        }

        TiedUpMod.LOGGER.info(
            "[PrisonerService] Extracted {} from cell {} by {}",
            captive.getName().getString(),
            cell.getId().toString().substring(0, 8),
            newCaptor.getEntity().getName().getString()
        );
        return true;
    }

    // ==================== RETURN TO CELL ====================

    /**
     * Return a prisoner to their cell.
     * WORKING → IMPRISONED + unleash + CellRegistryV2.assign.
     *
     * For: dogwalk return, labor return.
     * Note: If the entity is not tracked by PrisonerManager (e.g., damsels),
     * the PrisonerManager step is skipped — only unleash + CellRegistryV2 are handled.
     *
     * @param level Server level
     * @param captive The working entity
     * @param currentCaptor The current captor holding the leash
     * @param cell The cell to return to
     * @return true if return succeeded
     */
    public boolean returnToCell(
        ServerLevel level,
        LivingEntity captive,
        @Nullable ICaptor currentCaptor,
        CellDataV2 cell
    ) {
        UUID captiveId = captive.getUUID();
        PrisonerManager manager = PrisonerManager.get(level);
        CellRegistryV2 cellRegistry = CellRegistryV2.get(level);
        long gameTime = level.getGameTime();

        // 1. Unleash
        IBondageState kidnappedState = KidnappedHelper.getKidnappedState(
            captive
        );
        if (kidnappedState != null && kidnappedState.isCaptive()) {
            kidnappedState.free(false);
        }
        if (currentCaptor != null && kidnappedState != null) {
            currentCaptor.removeCaptive(kidnappedState, false);
        }

        // 2. PrisonerManager: WORKING → IMPRISONED (players only — damsels aren't tracked)
        if (captive instanceof Player) {
            boolean returned = manager.returnToCell(
                captiveId,
                cell.getId(),
                gameTime
            );
            if (!returned) {
                // State might not be WORKING (escape during return, etc.)
                // Force-repair: try capture+imprison path
                PrisonerRecord record = manager.getRecord(captiveId);
                PrisonerState currentState = record.getState();
                TiedUpMod.LOGGER.warn(
                    "[PrisonerService] returnToCell() transition failed for {} (state={}). Force-repairing.",
                    captiveId.toString().substring(0, 8),
                    currentState
                );

                if (currentState == PrisonerState.FREE) {
                    UUID captorId =
                        currentCaptor != null
                            ? currentCaptor.getEntity().getUUID()
                            : captiveId;
                    manager.capture(captiveId, captorId, gameTime);
                    UUID campId = record.getCampId();
                    manager.imprison(captiveId, campId, cell.getId(), gameTime);
                } else if (currentState == PrisonerState.CAPTURED) {
                    UUID campId = record.getCampId();
                    manager.imprison(captiveId, campId, cell.getId(), gameTime);
                }
                // If IMPRISONED already — no transition needed, just reassign cell
            }
        }

        // 3. Reassign to cell registry
        cellRegistry.assignPrisoner(cell.getId(), captiveId);

        TiedUpMod.LOGGER.info(
            "[PrisonerService] Returned {} to cell {}",
            captive.getName().getString(),
            cell.getId().toString().substring(0, 8)
        );
        return true;
    }

    // ==================== TRANSFER ====================

    /**
     * Transfer a captive from one captor to another.
     * PrisonerManager state unchanged.
     *
     * For: maid→buyer, kidnapper→kidnapper.
     *
     * @param level Server level
     * @param captive The captive entity
     * @param oldCaptor The current captor
     * @param newCaptor The new captor
     * @return true if transfer succeeded
     */
    public boolean transferCaptive(
        ServerLevel level,
        LivingEntity captive,
        ICaptor oldCaptor,
        ICaptor newCaptor
    ) {
        IBondageState kidnappedState = KidnappedHelper.getKidnappedState(
            captive
        );
        if (kidnappedState == null) {
            TiedUpMod.LOGGER.warn(
                "[PrisonerService] transferCaptive() no kidnapped state for {}",
                captive.getName().getString()
            );
            return false;
        }

        // 1. Remove from old captor
        oldCaptor.removeCaptive(kidnappedState, false);

        // 2. Clear leash
        kidnappedState.free(false);

        // 3. Re-leash to new captor
        boolean leashed;
        if (captive instanceof AbstractTiedUpNpc npc) {
            leashed = npc.forceCapturedBy(newCaptor);
        } else {
            leashed = kidnappedState.getCapturedBy(newCaptor);
        }

        if (!leashed) {
            // Rollback: try to re-leash to old captor
            TiedUpMod.LOGGER.warn(
                "[PrisonerService] transferCaptive() re-leash to new captor failed, attempting rollback"
            );
            if (captive instanceof AbstractTiedUpNpc npc) {
                npc.forceCapturedBy(oldCaptor);
            } else {
                kidnappedState.getCapturedBy(oldCaptor);
            }
            return false;
        }

        TiedUpMod.LOGGER.info(
            "[PrisonerService] Transferred {} from {} to {}",
            captive.getName().getString(),
            oldCaptor.getEntity().getName().getString(),
            newCaptor.getEntity().getName().getString()
        );
        return true;
    }

}