package com.tiedup.remake.entities.ai.kidnapper;

import com.tiedup.remake.cells.CellDataV2;
import com.tiedup.remake.cells.CellRegistryV2;
import com.tiedup.remake.core.TiedUpMod;
import com.tiedup.remake.dialogue.DialogueBridge;
import com.tiedup.remake.entities.AbstractTiedUpNpc;
import com.tiedup.remake.entities.EntityKidnapper;
import com.tiedup.remake.entities.ai.StuckDetector;
import com.tiedup.remake.entities.ai.WaypointNavigator;
import com.tiedup.remake.v2.bondage.datadriven.DataDrivenBondageItem;
import net.minecraft.resources.ResourceLocation;
import com.tiedup.remake.state.IBondageState;
import com.tiedup.remake.util.KidnappedHelper;
import com.tiedup.remake.util.KidnapperAIHelper;
import com.tiedup.remake.v2.BodyRegionV2;
import java.util.ArrayList;
import java.util.EnumSet;
import java.util.List;
import java.util.UUID;
import javax.annotation.Nullable;
import net.minecraft.core.BlockPos;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.server.level.ServerPlayer;
import net.minecraft.world.entity.ai.goal.Goal;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.level.Level;

/**
 * AI Goal for EntityKidnapper to take a prisoner for a walk.
 *
 * This goal:
 * 1. Randomly triggers when guarding cells with prisoners
 * 2. Teleports prisoner out of cell to kidnapper
 * 3. Changes their bind to DOGBINDER (pet play style)
 * 4. Leads them on a walk around the area
 * 5. Returns them to cell
 * 6. Changes bind back to normal
 *
 * Can also happen with Damsels (NPCs).
 *
 * Note: Only regular kidnappers can do this, not maids or traders.
 */
public class KidnapperWalkPrisonerGoal extends Goal {

    private final EntityKidnapper kidnapper;

    /** Walk states */
    private static final int STATE_IDLE = 0;
    private static final int STATE_PICKING_PRISONER = 1;
    private static final int STATE_GOING_TO_CELL = 2;
    private static final int STATE_WALKING = 3;
    private static final int STATE_RETURNING = 4;
    private static final int STATE_RECOVERING = 5;
    private static final int STATE_DONE = 6;

    /** Current state */
    private int walkState = STATE_IDLE;

    /** Current prisoner being walked */
    @Nullable
    private IBondageState walkedPrisoner;

    /** UUID of prisoner being walked (for recovery after escape) */
    @Nullable
    private UUID walkedPrisonerId;

    /** Original bind item (to restore after walk) */
    @Nullable
    private ItemStack originalBind;

    /** Cell the prisoner came from */
    @Nullable
    private CellDataV2 originCell;

    /** Waypoint navigator for going to/returning from cell */
    @Nullable
    private WaypointNavigator waypointNav;

    /** Current walk waypoint */
    @Nullable
    private BlockPos currentWaypoint;

    /** Ticks spent in current state */
    private int stateTicks;

    /** Walk duration (2-4 minutes) */
    private int walkDuration;

    /** Minimum walk duration (2 minutes) */
    private static final int MIN_WALK_DURATION = 2400;

    /** Maximum walk duration (4 minutes) */
    private static final int MAX_WALK_DURATION = 4800;

    /** Timeout for transition states (30 seconds) */
    private static final int TRANSITION_STATE_TIMEOUT = 600;

    /** Interval between waypoint changes (10-20 seconds) */
    private static final int WAYPOINT_INTERVAL_MIN = 200;
    private static final int WAYPOINT_INTERVAL_MAX = 400;

    /** Timer for waypoint changes */
    private int waypointTimer;

    /** Walk radius from cell */
    private static final double WALK_RADIUS = 30.0;

    /** Chance to trigger walk behavior (1/1000 per tick when guarding = ~5% per minute) */
    private static final int WALK_TRIGGER_CHANCE = 1000;

    /** Whether prisoner has been extracted from PrisonerManager */
    private boolean extracted = false;

    /** Cooldown after completing a walk (5 minutes) */
    private int walkCooldown = 0;
    private static final int WALK_COOLDOWN_DURATION = 6000;

    /** Stuck detection for vanilla navigation in GOING_TO_CELL and RETURNING states */
    private final StuckDetector stuckDetector = new StuckDetector(60, 3);

    public KidnapperWalkPrisonerGoal(EntityKidnapper kidnapper) {
        this.kidnapper = kidnapper;
        this.setFlags(EnumSet.of(Goal.Flag.MOVE, Goal.Flag.LOOK));
    }

    @Override
    public boolean canUse() {
        // Cooldown check
        if (this.walkCooldown > 0) {
            this.walkCooldown--;
            return false;
        }

        // Must not be tied up
        if (this.kidnapper.isTiedUp()) {
            return false;
        }

        // Must not already have a captive
        if (this.kidnapper.hasCaptives()) {
            return false;
        }

        // Must not have an active target
        if (this.kidnapper.getTarget() != null) {
            return false;
        }

        // Must be in a walkable state (GUARD, PATROL, or IDLE)
        KidnapperState state = this.kidnapper.getCurrentState();
        if (
            state != KidnapperState.GUARD &&
            state != KidnapperState.PATROL &&
            state != KidnapperState.IDLE
        ) {
            return false;
        }

        // Must not be waiting for job
        if (this.kidnapper.isWaitingForJobToBeCompleted()) {
            return false;
        }

        // Random chance trigger
        if (this.kidnapper.getRandom().nextInt(WALK_TRIGGER_CHANCE) != 0) {
            return false;
        }

        // Need nearby cells with prisoners
        List<CellDataV2> occupiedCells =
            this.kidnapper.getNearbyCellsWithPrisoners();
        if (occupiedCells.isEmpty()) {
            return false;
        }

        // Only one dogwalk per camp at a time
        if (isCampAlreadyDogwalking()) {
            return false;
        }

        return true;
    }

    /**
     * Check if any other kidnapper in this camp is already dogwalking.
     */
    private boolean isCampAlreadyDogwalking() {
        if (!(this.kidnapper.level() instanceof ServerLevel serverLevel)) {
            return false;
        }

        // Find nearby kidnappers
        for (EntityKidnapper other : serverLevel.getEntitiesOfClass(
            EntityKidnapper.class,
            this.kidnapper.getBoundingBox().inflate(100)
        )) {
            if (other == this.kidnapper) continue;

            // Check if this kidnapper is dogwalking
            if (other.isDogwalking()) {
                return true;
            }
        }

        return false;
    }

    @Override
    public boolean canContinueToUse() {
        if (this.walkState == STATE_DONE) {
            return false;
        }

        if (this.kidnapper.isTiedUp()) {
            return false;
        }

        // If prisoner escaped or died during walk
        if (
            this.walkedPrisoner != null &&
            !this.walkedPrisoner.asLivingEntity().isAlive()
        ) {
            return false;
        }

        return true;
    }

    @Override
    public void start() {
        this.walkState = STATE_PICKING_PRISONER;
        this.stateTicks = 0;
        this.extracted = false;
        this.stuckDetector.reset();
        this.walkDuration =
            MIN_WALK_DURATION +
            this.kidnapper.getRandom().nextInt(
                MAX_WALK_DURATION - MIN_WALK_DURATION
            );

        TiedUpMod.LOGGER.debug(
            "[KidnapperWalkPrisonerGoal] {} starting walk behavior",
            this.kidnapper.getNpcName()
        );
    }

    @Override
    public void stop() {
        // If walk was interrupted, restore prisoner to cell only if extracted
        if (this.walkedPrisoner != null && this.walkState != STATE_DONE) {
            if (this.extracted) {
                // Prisoner was extracted from cell, must return
                restorePrisonerToCell();
            }
            // else: not yet extracted (PICKING/GOING_TO_CELL), nothing to clean up
        }

        // Reset cooldown
        this.walkCooldown = WALK_COOLDOWN_DURATION;

        // Clear dogwalking flag
        this.kidnapper.setDogwalking(false);

        // Cleanup
        this.walkedPrisoner = null;
        this.walkedPrisonerId = null;
        this.originalBind = null;
        this.originCell = null;
        this.currentWaypoint = null;
        this.waypointNav = null;
        this.extracted = false;
        this.stuckDetector.reset();
        this.walkState = STATE_IDLE;
        this.kidnapper.getNavigation().stop();

        // Return to IDLE state
        if (this.kidnapper.getCurrentState() != KidnapperState.ALERT) {
            this.kidnapper.setCurrentState(KidnapperState.IDLE);
        }
    }

    @Override
    public void tick() {
        this.stateTicks++;

        // Update dogwalking flag on kidnapper
        this.kidnapper.setDogwalking(
            this.walkState != STATE_IDLE && this.walkState != STATE_DONE
        );

        switch (this.walkState) {
            case STATE_PICKING_PRISONER -> tickPickPrisoner();
            case STATE_GOING_TO_CELL -> tickGoingToCell();
            case STATE_WALKING -> tickWalking();
            case STATE_RETURNING -> tickReturning();
            case STATE_RECOVERING -> tickRecovering();
        }
    }

    /**
     * Pick a prisoner from a cell to walk.
     * This state selects the cell and prisoner, then transitions to GOING_TO_CELL.
     */
    private void tickPickPrisoner() {
        if (!(this.kidnapper.level() instanceof ServerLevel serverLevel)) {
            this.walkState = STATE_DONE;
            return;
        }

        // Find a cell with prisoners
        List<CellDataV2> occupiedCells =
            this.kidnapper.getNearbyCellsWithPrisoners();
        if (occupiedCells.isEmpty()) {
            this.walkState = STATE_DONE;
            return;
        }

        // Pick a random cell
        CellDataV2 cell = occupiedCells.get(
            this.kidnapper.getRandom().nextInt(occupiedCells.size())
        );
        List<UUID> prisonerIds = cell.getPrisonerIds();
        if (prisonerIds.isEmpty()) {
            this.walkState = STATE_DONE;
            return;
        }

        // Pick a random prisoner (player or damsel)
        UUID prisonerId = prisonerIds.get(
            this.kidnapper.getRandom().nextInt(prisonerIds.size())
        );

        // Try to get the prisoner entity
        IBondageState prisoner = findPrisoner(serverLevel, prisonerId);
        if (prisoner == null) {
            this.walkState = STATE_DONE;
            return;
        }

        // Check if prisoner is tied up
        if (!prisoner.isTiedUp()) {
            this.walkState = STATE_DONE;
            return;
        }

        // Store selection - actual pickup happens when we reach the cell
        this.walkedPrisoner = prisoner;
        this.walkedPrisonerId = prisonerId;
        this.originCell = cell;
        this.originalBind = prisoner.getEquipment(BodyRegionV2.ARMS).copy();

        TiedUpMod.LOGGER.debug(
            "[KidnapperWalkPrisonerGoal] {} going to cell to pick up {}",
            this.kidnapper.getNpcName(),
            prisoner.getKidnappedName()
        );

        // Set up waypoint navigation if cell has waypoints (forward to cell)
        if (!cell.getPathWaypoints().isEmpty()) {
            List<BlockPos> path = new ArrayList<>(cell.getPathWaypoints());
            BlockPos deliveryPoint = getDeliveryOrSpawnPoint(cell);
            path.add(deliveryPoint); // Add delivery as final waypoint
            this.waypointNav = new WaypointNavigator(this.kidnapper, path, 0.8);
            TiedUpMod.LOGGER.debug(
                "[KidnapperWalkPrisonerGoal] Using {} waypoints to reach cell",
                path.size()
            );
        } else {
            this.waypointNav = null;
        }

        // Transition to going to cell
        this.walkState = STATE_GOING_TO_CELL;
        this.stateTicks = 0;
        this.stuckDetector.reset();
    }

    /**
     * Walk to the cell to pick up the prisoner.
     */
    private void tickGoingToCell() {
        // TIMEOUT: Prevent infinite pathfinding if stuck
        if (this.stateTicks > TRANSITION_STATE_TIMEOUT) {
            TiedUpMod.LOGGER.warn(
                "[KidnapperWalkPrisonerGoal] {} timed out going to cell, aborting walk",
                this.kidnapper.getNpcName()
            );
            this.walkState = STATE_DONE;
            return;
        }

        if (this.originCell == null || this.walkedPrisoner == null) {
            this.walkState = STATE_DONE;
            return;
        }

        // Check if prisoner is still valid
        if (
            !this.walkedPrisoner.asLivingEntity().isAlive() ||
            !this.walkedPrisoner.isTiedUp()
        ) {
            this.walkState = STATE_DONE;
            return;
        }

        // WAYPOINT NAVIGATION MODE
        if (this.waypointNav != null && this.waypointNav.hasWaypoints()) {
            this.waypointNav.tick();

            if (this.waypointNav.isComplete()) {
                // Reached cell - pick up the prisoner
                pickUpPrisoner();
                return;
            }

            // Navigate to current waypoint
            if (this.stateTicks % 20 == 0) {
                this.waypointNav.navigateToCurrentWaypoint();
            }
            return;
        }

        // VANILLA NAVIGATION MODE
        BlockPos deliveryPoint = getDeliveryOrSpawnPoint(this.originCell);
        double distToCell = this.kidnapper.distanceToSqr(
            deliveryPoint.getX() + 0.5,
            deliveryPoint.getY(),
            deliveryPoint.getZ() + 0.5
        );

        // Arrived at cell - pick up the prisoner
        if (distToCell <= 9.0) {
            pickUpPrisoner();
            return;
        }

        // Update stuck detection
        this.stuckDetector.update(this.kidnapper);

        if (this.stuckDetector.isStuck()) {
            if (this.stuckDetector.shouldTeleport()) {
                // Teleport to delivery point
                this.kidnapper.teleportTo(
                    deliveryPoint.getX() + 0.5,
                    deliveryPoint.getY(),
                    deliveryPoint.getZ() + 0.5
                );
                TiedUpMod.LOGGER.debug(
                    "[KidnapperWalkPrisonerGoal] {} teleported to cell delivery point to unblock",
                    this.kidnapper.getNpcName()
                );
                this.stuckDetector.reset();
                pickUpPrisoner();
            } else if (
                this.kidnapper.level() instanceof ServerLevel serverLevel
            ) {
                BlockPos detour = this.stuckDetector.tryDetour(
                    serverLevel,
                    this.kidnapper.blockPosition(),
                    5
                );
                if (detour != null) {
                    this.kidnapper.getNavigation().moveTo(
                        detour.getX() + 0.5,
                        detour.getY(),
                        detour.getZ() + 0.5,
                        0.8
                    );
                }
            }
            return;
        }

        // Skip normal navigation if in detour
        if (this.stuckDetector.isInDetour()) {
            return;
        }

        // Move towards cell delivery point
        if (this.stateTicks % 20 == 0) {
            boolean pathFound = this.kidnapper.getNavigation().moveTo(
                deliveryPoint.getX() + 0.5,
                deliveryPoint.getY(),
                deliveryPoint.getZ() + 0.5,
                0.8
            );
            if (!pathFound) {
                this.stuckDetector.onPathFailed();
            }
        }
    }

    /**
     * Actually pick up the prisoner once at the cell.
     *
     * Uses PrisonerService.extractFromCell() for atomic state management:
     * 1. PrisonerService.extractFromCell() — PrisonerManager + CellRegistry + leash (with rollback)
     * 2. Teleport prisoner to kidnapper
     * 3. Change bind to DOGBINDER
     */
    private void pickUpPrisoner() {
        if (this.walkedPrisoner == null || this.walkedPrisonerId == null) {
            this.walkState = STATE_DONE;
            return;
        }

        if (!(this.kidnapper.level() instanceof ServerLevel serverLevel)) {
            this.walkState = STATE_DONE;
            return;
        }

        if (this.originCell == null) {
            this.walkState = STATE_DONE;
            return;
        }

        net.minecraft.world.entity.LivingEntity prisonerEntity =
            this.walkedPrisoner.asLivingEntity();

        // 1. Extract from cell via PrisonerService (PrisonerManager + CellRegistry + leash, atomic)
        boolean success =
            com.tiedup.remake.prison.service.PrisonerService.get().extractFromCell(
                serverLevel,
                prisonerEntity,
                this.kidnapper,
                this.originCell
            );

        if (!success) {
            TiedUpMod.LOGGER.warn(
                "[KidnapperWalkPrisonerGoal] {} failed to extract {}, aborting walk",
                this.kidnapper.getNpcName(),
                this.walkedPrisoner.getKidnappedName()
            );
            this.walkState = STATE_DONE;
            return;
        }

        this.extracted = true;

        // 2. Teleport prisoner to kidnapper
        prisonerEntity.teleportTo(
            this.kidnapper.getX(),
            this.kidnapper.getY(),
            this.kidnapper.getZ()
        );

        // 3. Change bind to DOGBINDER
        ItemStack dogBinder = DataDrivenBondageItem.createStack(ResourceLocation.fromNamespaceAndPath("tiedup", "dogbinder"));
        this.walkedPrisoner.equip(BodyRegionV2.ARMS, dogBinder);

        TiedUpMod.LOGGER.debug(
            "[KidnapperWalkPrisonerGoal] {} taking {} for a walk",
            this.kidnapper.getNpcName(),
            this.walkedPrisoner.getKidnappedName()
        );

        // Dialogue: Start walk
        if (prisonerEntity instanceof ServerPlayer player) {
            DialogueBridge.talkTo(this.kidnapper, player, "dogwalk.start");
        }

        this.walkState = STATE_WALKING;
        this.stateTicks = 0;
        this.waypointTimer = 0;
        this.currentWaypoint = calculateWalkWaypoint();
    }

    /**
     * Walk the prisoner around.
     */
    private void tickWalking() {
        if (this.walkedPrisoner == null) {
            this.walkState = STATE_DONE;
            return;
        }

        // Check if captive was lost (rope broke, fight back, etc.)
        if (!this.kidnapper.hasCaptives()) {
            TiedUpMod.LOGGER.debug(
                "[KidnapperWalkPrisonerGoal] {} lost captive during walk, switching to recovery",
                this.kidnapper.getNpcName()
            );
            this.walkState = STATE_RECOVERING;
            this.stateTicks = 0;
            return;
        }

        this.waypointTimer++;

        // Change waypoint periodically
        int waypointInterval =
            WAYPOINT_INTERVAL_MIN +
            this.kidnapper.getRandom().nextInt(
                WAYPOINT_INTERVAL_MAX - WAYPOINT_INTERVAL_MIN
            );
        if (
            this.waypointTimer >= waypointInterval ||
            this.currentWaypoint == null
        ) {
            this.waypointTimer = 0;
            this.currentWaypoint = calculateWalkWaypoint();
        }

        // Occasional dialogue during walk (every ~20 seconds)
        if (
            this.stateTicks % 400 == 0 &&
            this.kidnapper.getRandom().nextFloat() < 0.5f
        ) {
            if (
                this.walkedPrisoner.asLivingEntity() instanceof
                    ServerPlayer player
            ) {
                DialogueBridge.talkTo(this.kidnapper, player, "dogwalk.during");
            }
        }

        // Move towards waypoint
        if (this.currentWaypoint != null) {
            double distSq = this.kidnapper.distanceToSqr(
                this.currentWaypoint.getX() + 0.5,
                this.currentWaypoint.getY(),
                this.currentWaypoint.getZ() + 0.5
            );

            if (distSq > 4.0) {
                if (this.stateTicks % 20 == 0) {
                    this.kidnapper.getNavigation().moveTo(
                        this.currentWaypoint.getX() + 0.5,
                        this.currentWaypoint.getY(),
                        this.currentWaypoint.getZ() + 0.5,
                        0.7 // Moderate walking speed
                    );
                }
            }
        }

        // Look at prisoner occasionally
        if (this.stateTicks % 40 == 0) {
            this.kidnapper.getLookControl().setLookAt(
                this.walkedPrisoner.asLivingEntity(),
                30.0F,
                30.0F
            );
        }

        // Check if walk duration is complete
        if (this.stateTicks >= this.walkDuration) {
            TiedUpMod.LOGGER.debug(
                "[KidnapperWalkPrisonerGoal] {} finished walking {}, returning to cell",
                this.kidnapper.getNpcName(),
                this.walkedPrisoner.getKidnappedName()
            );

            // Set up waypoint navigation for return to cell (forward direction - going TO cell)
            if (
                this.originCell != null &&
                !this.originCell.getPathWaypoints().isEmpty()
            ) {
                List<BlockPos> path = new ArrayList<>(
                    this.originCell.getPathWaypoints()
                );
                BlockPos deliveryPoint = getDeliveryOrSpawnPoint(
                    this.originCell
                );
                path.add(deliveryPoint); // Add delivery as final waypoint
                this.waypointNav = new WaypointNavigator(
                    this.kidnapper,
                    path,
                    0.8
                );
                TiedUpMod.LOGGER.debug(
                    "[KidnapperWalkPrisonerGoal] Using {} waypoints to return to cell",
                    path.size()
                );
            } else {
                this.waypointNav = null;
            }

            this.walkState = STATE_RETURNING;
            this.stateTicks = 0;
            this.stuckDetector.reset();
        }
    }

    /**
     * Recover the escaped prisoner and resume walk.
     */
    private void tickRecovering() {
        if (this.walkedPrisonerId == null) {
            this.walkState = STATE_DONE;
            return;
        }

        if (!(this.kidnapper.level() instanceof ServerLevel serverLevel)) {
            this.walkState = STATE_DONE;
            return;
        }

        // Find the prisoner
        IBondageState prisoner = findPrisoner(
            serverLevel,
            this.walkedPrisonerId
        );
        if (prisoner == null || !prisoner.asLivingEntity().isAlive()) {
            // Prisoner died or disappeared - abort
            TiedUpMod.LOGGER.debug(
                "[KidnapperWalkPrisonerGoal] {} lost prisoner during recovery, aborting walk",
                this.kidnapper.getNpcName()
            );
            this.walkState = STATE_DONE;
            return;
        }

        net.minecraft.world.entity.LivingEntity prisonerEntity =
            prisoner.asLivingEntity();

        // Move towards the prisoner
        double distSq = this.kidnapper.distanceToSqr(prisonerEntity);
        if (distSq > 4.0) {
            if (this.stateTicks % 10 == 0) {
                this.kidnapper.getNavigation().moveTo(prisonerEntity, 1.0); // Run to catch up
            }
            this.kidnapper.getLookControl().setLookAt(
                prisonerEntity,
                30.0F,
                30.0F
            );
            return;
        }

        // Close enough - try to recapture
        // Make sure they still have the dogbinder on
        ItemStack currentBind = prisoner.getEquipment(BodyRegionV2.ARMS);
        if (currentBind.isEmpty() || !prisoner.isTiedUp()) {
            // They freed themselves - put dogbinder back on
            ItemStack dogBinder = new ItemStack(
                DataDrivenBondageItem.createStack(ResourceLocation.fromNamespaceAndPath("tiedup", "dogbinder")).getItem()
            );
            prisoner.equip(BodyRegionV2.ARMS, dogBinder);
        }

        // Recapture — use forceCapturedBy for NPCs (bypasses canCapture)
        boolean recaptured;
        if (prisoner.asLivingEntity() instanceof AbstractTiedUpNpc npc) {
            recaptured = npc.forceCapturedBy(this.kidnapper);
        } else {
            recaptured = prisoner.getCapturedBy(this.kidnapper);
        }

        if (recaptured) {
            TiedUpMod.LOGGER.debug(
                "[KidnapperWalkPrisonerGoal] {} recaptured {}, resuming walk",
                this.kidnapper.getNpcName(),
                prisoner.getKidnappedName()
            );
            this.walkedPrisoner = prisoner;
            this.walkState = STATE_WALKING;
            this.stateTicks = 0;
        }
    }

    /**
     * Return prisoner to their cell.
     */
    private void tickReturning() {
        // TIMEOUT: Prevent infinite pathfinding if stuck
        if (this.stateTicks > TRANSITION_STATE_TIMEOUT) {
            TiedUpMod.LOGGER.warn(
                "[KidnapperWalkPrisonerGoal] {} timed out returning to cell, forcing restore via teleport",
                this.kidnapper.getNpcName()
            );
            restorePrisonerToCell(); // Force restore via teleport fallback
            this.walkState = STATE_DONE;
            return;
        }

        if (this.walkedPrisoner == null || this.originCell == null) {
            this.walkState = STATE_DONE;
            return;
        }

        // Check if captive was lost during return
        if (!this.kidnapper.hasCaptives()) {
            TiedUpMod.LOGGER.debug(
                "[KidnapperWalkPrisonerGoal] {} lost captive during return, switching to recovery",
                this.kidnapper.getNpcName()
            );
            this.walkState = STATE_RECOVERING;
            this.stateTicks = 0;
            return;
        }

        // WAYPOINT NAVIGATION MODE
        if (this.waypointNav != null && this.waypointNav.hasWaypoints()) {
            this.waypointNav.tick();

            if (this.waypointNav.isComplete()) {
                // Arrived at cell - restore prisoner
                restorePrisonerToCell();
                this.walkState = STATE_DONE;
                return;
            }

            // Navigate to current waypoint
            if (this.stateTicks % 20 == 0) {
                this.waypointNav.navigateToCurrentWaypoint();
            }
            return;
        }

        // VANILLA NAVIGATION MODE
        BlockPos deliveryPoint = getDeliveryOrSpawnPoint(this.originCell);
        double distToCell = this.kidnapper.distanceToSqr(
            deliveryPoint.getX() + 0.5,
            deliveryPoint.getY(),
            deliveryPoint.getZ() + 0.5
        );

        // Arrived at cell - restore prisoner
        if (distToCell <= 9.0) {
            restorePrisonerToCell();
            this.walkState = STATE_DONE;
            return;
        }

        // Update stuck detection
        this.stuckDetector.update(this.kidnapper);

        if (this.stuckDetector.isStuck()) {
            if (this.stuckDetector.shouldTeleport()) {
                // Teleport to delivery point and restore prisoner
                this.kidnapper.teleportTo(
                    deliveryPoint.getX() + 0.5,
                    deliveryPoint.getY(),
                    deliveryPoint.getZ() + 0.5
                );
                TiedUpMod.LOGGER.debug(
                    "[KidnapperWalkPrisonerGoal] {} teleported to delivery point for return",
                    this.kidnapper.getNpcName()
                );
                this.stuckDetector.reset();
                restorePrisonerToCell();
                this.walkState = STATE_DONE;
            } else if (
                this.kidnapper.level() instanceof ServerLevel serverLevel
            ) {
                BlockPos detour = this.stuckDetector.tryDetour(
                    serverLevel,
                    this.kidnapper.blockPosition(),
                    5
                );
                if (detour != null) {
                    this.kidnapper.getNavigation().moveTo(
                        detour.getX() + 0.5,
                        detour.getY(),
                        detour.getZ() + 0.5,
                        0.8
                    );
                }
            }
            return;
        }

        // Skip normal navigation if in detour
        if (this.stuckDetector.isInDetour()) {
            return;
        }

        // Move towards cell delivery point
        if (this.stateTicks % 20 == 0) {
            boolean pathFound = this.kidnapper.getNavigation().moveTo(
                deliveryPoint.getX() + 0.5,
                deliveryPoint.getY(),
                deliveryPoint.getZ() + 0.5,
                0.8
            );
            if (!pathFound) {
                this.stuckDetector.onPathFailed();
            }
        }
    }

    /**
     * Restore prisoner to their cell with original bind.
     * Uses PrisonerService.returnToCell() for atomic state management.
     * Uses findGroundPosition() to ensure safe teleport location.
     * Handles dead prisoners: still fixes PrisonerManager/CellRegistry state.
     */
    private void restorePrisonerToCell() {
        if (this.walkedPrisoner == null) return;

        if (
            !(this.kidnapper.level() instanceof ServerLevel serverLevel)
        ) return;

        net.minecraft.world.entity.LivingEntity prisonerEntity =
            this.walkedPrisoner.asLivingEntity();
        boolean isAlive = prisonerEntity.isAlive();

        // 1. Restore original bind (only if alive)
        if (
            isAlive && this.originalBind != null && !this.originalBind.isEmpty()
        ) {
            this.walkedPrisoner.equip(BodyRegionV2.ARMS, this.originalBind);
        }

        // 2. Teleport to cell spawn point (only if alive)
        if (isAlive && this.originCell != null) {
            BlockPos spawn =
                this.originCell.getSpawnPoint() != null
                    ? this.originCell.getSpawnPoint()
                    : this.originCell.getCorePos().above();
            BlockPos safeSpawn = findGroundPosition(serverLevel, spawn);
            prisonerEntity.teleportTo(
                safeSpawn.getX() + 0.5,
                safeSpawn.getY(),
                safeSpawn.getZ() + 0.5
            );
        }

        // 3. Return to cell via PrisonerService (unleash + PrisonerManager + CellRegistry, with force-repair)
        if (this.extracted && this.originCell != null) {
            com.tiedup.remake.prison.service.PrisonerService.get().returnToCell(
                serverLevel,
                prisonerEntity,
                this.kidnapper,
                this.originCell
            );
            this.extracted = false;
        } else if (this.originCell != null && this.walkedPrisonerId != null) {
            // Not extracted (damsel or PrisonerManager wasn't involved) — just unleash + reassign cell
            if (isAlive && this.walkedPrisoner.isCaptive()) {
                this.walkedPrisoner.free(false);
            }
            CellRegistryV2 registry = CellRegistryV2.get(serverLevel);
            registry.assignPrisoner(
                this.originCell.getId(),
                this.walkedPrisonerId
            );
        }

        // 4. Dialogue: Return to cell (only if alive)
        if (isAlive && prisonerEntity instanceof ServerPlayer player) {
            DialogueBridge.talkTo(this.kidnapper, player, "dogwalk.return");
        }

        TiedUpMod.LOGGER.debug(
            "[KidnapperWalkPrisonerGoal] {} returned {} to cell{}",
            this.kidnapper.getNpcName(),
            this.walkedPrisoner.getKidnappedName(),
            isAlive ? "" : " (dead - state cleanup only)"
        );
    }

    /**
     * Find a prisoner by UUID - can be player or NPC.
     */
    @Nullable
    private IBondageState findPrisoner(ServerLevel level, UUID prisonerId) {
        // Try player first
        ServerPlayer player = level
            .getServer()
            .getPlayerList()
            .getPlayer(prisonerId);
        if (player != null && player.isAlive()) {
            return KidnappedHelper.getKidnappedState(player);
        }

        // Try NPC
        net.minecraft.world.entity.Entity entity = level.getEntity(prisonerId);
        if (entity instanceof AbstractTiedUpNpc npc && npc.isAlive()) {
            return npc;
        }

        return null;
    }

    /**
     * Calculate a random walk waypoint around the origin cell.
     * Uses findGroundPosition() to ensure waypoints are on solid ground.
     */
    private BlockPos calculateWalkWaypoint() {
        BlockPos center;
        if (this.originCell != null) {
            center = this.originCell.getCorePos();
        } else {
            center = this.kidnapper.blockPosition();
        }

        // Random angle
        double angle = this.kidnapper.getRandom().nextDouble() * Math.PI * 2;

        // Random distance
        double distance =
            10 + this.kidnapper.getRandom().nextDouble() * (WALK_RADIUS - 10);

        int offsetX = (int) (Math.cos(angle) * distance);
        int offsetZ = (int) (Math.sin(angle) * distance);

        // Calculate raw position and find actual ground
        BlockPos rawPos = center.offset(offsetX, 0, offsetZ);
        return findGroundPosition(this.kidnapper.level(), rawPos);
    }

    /**
     * Find a safe ground position at or near the given XZ coordinates.
     * Searches down first (max 10 blocks), then up if not found.
     *
     * @param level The level to search in
     * @param pos   The starting position
     * @return A position on solid ground, or the original position if none found
     */
    private BlockPos findGroundPosition(Level level, BlockPos pos) {
        BlockPos.MutableBlockPos mutable = pos.mutable();

        // Search downward first (max 10 blocks)
        for (int y = 0; y > -10; y--) {
            mutable.setY(pos.getY() + y);
            if (
                level.getBlockState(mutable).isSolidRender(level, mutable) &&
                level.getBlockState(mutable.above()).isAir()
            ) {
                return mutable.above().immutable();
            }
        }

        // Search upward if ground not found below (max 10 blocks)
        for (int y = 1; y < 10; y++) {
            mutable.setY(pos.getY() + y);
            if (
                level
                    .getBlockState(mutable.below())
                    .isSolidRender(level, mutable.below()) &&
                level.getBlockState(mutable).isAir()
            ) {
                return mutable.immutable();
            }
        }

        // Fallback: return original position
        return pos;
    }

    /**
     * Get the delivery point for a cell, or fall back to spawn point.
     * Uses KidnapperAIHelper for shared cell navigation logic.
     */
    private BlockPos getDeliveryOrSpawnPoint(CellDataV2 cell) {
        if (cell == null) {
            return this.kidnapper.blockPosition();
        }
        return KidnapperAIHelper.getDeliveryOrSpawnPoint(
            cell,
            this.kidnapper.level()
        );
    }
}
