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

import com.tiedup.remake.entities.EntityDamsel;
import com.tiedup.remake.personality.NpcCommand;
import com.tiedup.remake.personality.PersonalityState;
import java.util.*;
import javax.annotation.Nullable;
import net.minecraft.core.BlockPos;
import net.minecraft.sounds.SoundEvents;
import net.minecraft.sounds.SoundSource;
import net.minecraft.world.Container;
import net.minecraft.world.entity.ai.goal.Goal;
import net.minecraft.world.entity.player.Player;
import net.minecraft.world.item.Item;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.level.Level;
import net.minecraft.world.level.block.entity.ChestBlockEntity;

/**
 * AI Goal: NPC sorts items by distributing them to matching chests.
 *
 * <p>Takes items from source chest (commandTarget) and distributes them
 * to nearby chests (16 block radius) that already contain the same item.
 * If no matching chest is found, the item stays in the source.</p>
 *
 * <p><b>Behavior:</b></p>
 * <ul>
 *   <li>Build content map of all chests in zone</li>
 *   <li>Take an item from source chest</li>
 *   <li>Find a destination chest containing the same item</li>
 *   <li>Move to destination and store it</li>
 *   <li>Repeat while source has sortable items</li>
 * </ul>
 */
public class NpcSortCommandGoal extends Goal {

    private final EntityDamsel npc;
    private final double speedModifier;

    /** Radius to search for destination chests */
    private static final int CHEST_SEARCH_RADIUS = 16;

    private static final double INTERACT_DISTANCE = 2.0;
    private static final int XP_INTERVAL = 1200;
    private static final int IDLE_RESCAN_INTERVAL = 100;
    private static final int ACTION_DELAY = 20;
    private static final int CACHE_REFRESH_INTERVAL = 200;
    private static final float WORK_INTENSITY = 0.6f;

    private enum SortPhase {
        BUILDING_MAP,
        MOVING_TO_SOURCE,
        COLLECTING,
        FINDING_DEST,
        MOVING_TO_DEST,
        STORING,
        IDLE,
    }

    private SortPhase phase;

    @Nullable
    private BlockPos sourceChest;

    @Nullable
    private BlockPos destChest;

    @Nullable
    private Player master;

    @Nullable
    private ItemStack heldItem;

    /** Cache: maps Item -> list of chest positions containing that item */
    private Map<Item, List<BlockPos>> chestContentMap = new HashMap<>();

    /** Ticks since last cache refresh */
    private int cacheAge;

    private int pathRecalcTimer;
    private int xpTimer;
    private int idleTimer;
    private int actionTimer;

    public NpcSortCommandGoal(EntityDamsel npc, double speedModifier) {
        this.npc = npc;
        this.speedModifier = speedModifier;
        this.setFlags(EnumSet.of(Goal.Flag.MOVE, Goal.Flag.LOOK));
    }

    @Override
    public boolean canUse() {
        if (npc.getActiveCommand() != NpcCommand.SORT) {
            return false;
        }
        if (!npc.hasCollar()) {
            return false;
        }

        this.master = NpcGoalHelper.findCommandingPlayer(npc);
        return this.master != null;
    }

    @Override
    public boolean canContinueToUse() {
        if (npc.getActiveCommand() != NpcCommand.SORT) {
            return false;
        }
        if (!npc.hasCollar()) {
            return false;
        }
        return true;
    }

    @Override
    public void start() {
        NpcGoalHelper.ensureOutsideCell(npc);
        var state = npc.getPersonalityState();
        if (state == null || state.getCommandTarget() == null) {
            npc.cancelCommand();
            return;
        }

        this.sourceChest = state.getCommandTarget();
        this.destChest = null;
        this.phase = SortPhase.BUILDING_MAP;
        this.chestContentMap.clear();
        this.cacheAge = 0;
        this.pathRecalcTimer = 0;
        this.xpTimer = 0;
        this.idleTimer = 0;
        this.actionTimer = 0;
        this.heldItem = null;
    }

    @Override
    public void stop() {
        // Return held item to source if possible
        if (heldItem != null && !heldItem.isEmpty() && sourceChest != null) {
            Container container = NpcGoalHelper.getChestContainer(
                npc.level(),
                sourceChest
            );
            if (container != null) {
                NpcGoalHelper.insertIntoContainer(container, heldItem);
            }
            if (heldItem != null && !heldItem.isEmpty()) {
                npc.spawnAtLocation(heldItem);
            }
        }

        this.sourceChest = null;
        this.destChest = null;
        this.master = null;
        this.heldItem = null;
        this.chestContentMap.clear();
        npc.getNavigation().stop();
    }

    @Override
    public void tick() {
        // Award XP periodically
        if (++this.xpTimer >= XP_INTERVAL) {
            this.xpTimer = 0;
            awardXP(2);
            NpcGoalHelper.incrementJobExperience(npc, NpcCommand.SORT);
        }

        // Age the cache
        this.cacheAge++;

        switch (this.phase) {
            case BUILDING_MAP -> tickBuildingMap();
            case MOVING_TO_SOURCE -> tickMovingToSource();
            case COLLECTING -> tickCollecting();
            case FINDING_DEST -> tickFindingDest();
            case MOVING_TO_DEST -> tickMovingToDest();
            case STORING -> tickStoring();
            case IDLE -> tickIdle();
        }
    }

    private void tickBuildingMap() {
        buildChestContentMap();
        this.cacheAge = 0;
        this.phase = SortPhase.MOVING_TO_SOURCE;
        this.pathRecalcTimer = 0;
    }

    private void tickMovingToSource() {
        if (this.sourceChest == null) {
            this.phase = SortPhase.IDLE;
            return;
        }

        if (--this.pathRecalcTimer <= 0) {
            this.pathRecalcTimer = 10;
            npc
                .getNavigation()
                .moveTo(
                    sourceChest.getX() + 0.5,
                    sourceChest.getY(),
                    sourceChest.getZ() + 0.5,
                    speedModifier
                );
        }

        double dist = npc.blockPosition().distSqr(sourceChest);
        if (dist <= INTERACT_DISTANCE * INTERACT_DISTANCE) {
            npc.getNavigation().stop();
            this.phase = SortPhase.COLLECTING;
            this.actionTimer = ACTION_DELAY;
        }
    }

    private void tickCollecting() {
        if (sourceChest == null) {
            this.phase = SortPhase.IDLE;
            return;
        }

        if (--this.actionTimer > 0) {
            return;
        }

        // Refresh cache if stale
        if (this.cacheAge >= CACHE_REFRESH_INTERVAL) {
            this.phase = SortPhase.BUILDING_MAP;
            return;
        }

        Level level = npc.level();
        Container container = NpcGoalHelper.getChestContainer(
            level,
            sourceChest
        );

        if (container == null) {
            this.phase = SortPhase.IDLE;
            return;
        }

        // Find first item that has a matching destination chest
        for (int i = 0; i < container.getContainerSize(); i++) {
            ItemStack stack = container.getItem(i);
            if (stack.isEmpty()) continue;

            Item item = stack.getItem();
            List<BlockPos> destinations = chestContentMap.get(item);
            if (destinations != null && !destinations.isEmpty()) {
                // Take the stack
                this.heldItem = stack.copy();
                container.setItem(i, ItemStack.EMPTY);
                container.setChanged();

                level.playSound(
                    null,
                    sourceChest,
                    SoundEvents.CHEST_OPEN,
                    SoundSource.BLOCKS,
                    0.5f,
                    1.0f
                );

                // Drain rest
                PersonalityState personalityState = npc.getPersonalityState();
                if (personalityState != null) {
                    personalityState.getNeeds().drainFromWork(WORK_INTENSITY);
                }

                this.phase = SortPhase.FINDING_DEST;
                return;
            }
        }

        // No sortable items found
        this.phase = SortPhase.IDLE;
        this.idleTimer = 0;
    }

    private void tickFindingDest() {
        if (heldItem == null || heldItem.isEmpty()) {
            this.phase = SortPhase.MOVING_TO_SOURCE;
            return;
        }

        Item item = heldItem.getItem();
        List<BlockPos> destinations = chestContentMap.get(item);

        if (destinations == null || destinations.isEmpty()) {
            // No destination — return to source and put it back
            returnItemToSource();
            this.phase = SortPhase.MOVING_TO_SOURCE;
            this.pathRecalcTimer = 0;
            return;
        }

        // Pick the best destination (most items of same type = best consolidation)
        BlockPos best = findBestDestination(destinations, item);
        if (best == null) {
            returnItemToSource();
            this.phase = SortPhase.MOVING_TO_SOURCE;
            this.pathRecalcTimer = 0;
            return;
        }

        this.destChest = best;
        this.phase = SortPhase.MOVING_TO_DEST;
        this.pathRecalcTimer = 0;
    }

    private void tickMovingToDest() {
        if (this.destChest == null) {
            this.phase = SortPhase.FINDING_DEST;
            return;
        }

        if (--this.pathRecalcTimer <= 0) {
            this.pathRecalcTimer = 10;
            npc
                .getNavigation()
                .moveTo(
                    destChest.getX() + 0.5,
                    destChest.getY(),
                    destChest.getZ() + 0.5,
                    speedModifier
                );
        }

        double dist = npc.blockPosition().distSqr(destChest);
        if (dist <= INTERACT_DISTANCE * INTERACT_DISTANCE) {
            npc.getNavigation().stop();
            this.phase = SortPhase.STORING;
            this.actionTimer = ACTION_DELAY;
        }
    }

    private void tickStoring() {
        if (destChest == null || heldItem == null || heldItem.isEmpty()) {
            this.heldItem = null;
            this.destChest = null;
            this.phase = SortPhase.MOVING_TO_SOURCE;
            this.pathRecalcTimer = 0;
            return;
        }

        if (--this.actionTimer > 0) {
            return;
        }

        Level level = npc.level();
        Container container = NpcGoalHelper.getChestContainer(level, destChest);

        if (container != null) {
            ItemStack remaining = NpcGoalHelper.insertIntoContainer(
                container,
                heldItem
            );

            if (!remaining.isEmpty()) {
                // Couldn't store all — hold remainder for next trip
                this.heldItem = remaining;
            } else {
                this.heldItem = null;
            }

            level.playSound(
                null,
                destChest,
                SoundEvents.CHEST_CLOSE,
                SoundSource.BLOCKS,
                0.5f,
                1.0f
            );
        }

        awardXP(1);

        this.destChest = null;

        if (this.heldItem != null && !this.heldItem.isEmpty()) {
            // Still have items, find another dest or return to source
            this.phase = SortPhase.FINDING_DEST;
        } else {
            // Go back for more
            this.phase = SortPhase.MOVING_TO_SOURCE;
            this.pathRecalcTimer = 0;
        }
    }

    private void tickIdle() {
        if (++this.idleTimer >= IDLE_RESCAN_INTERVAL) {
            this.idleTimer = 0;

            // Rebuild cache and try again
            this.phase = SortPhase.BUILDING_MAP;
        }
    }

    // ========== Content Map ==========

    /**
     * Scan all chests within CHEST_SEARCH_RADIUS of the source chest,
     * build an index of which items are in which chests.
     * The source chest itself is excluded from destinations.
     */
    private void buildChestContentMap() {
        this.chestContentMap.clear();

        if (sourceChest == null) return;

        Level level = npc.level();
        BlockPos.MutableBlockPos mutable = new BlockPos.MutableBlockPos();

        for (int x = -CHEST_SEARCH_RADIUS; x <= CHEST_SEARCH_RADIUS; x++) {
            for (int z = -CHEST_SEARCH_RADIUS; z <= CHEST_SEARCH_RADIUS; z++) {
                for (int y = -2; y <= 2; y++) {
                    mutable.set(
                        sourceChest.getX() + x,
                        sourceChest.getY() + y,
                        sourceChest.getZ() + z
                    );

                    // Skip the source chest itself
                    if (mutable.equals(sourceChest)) continue;

                    if (
                        !(level.getBlockEntity(mutable) instanceof
                                ChestBlockEntity)
                    ) {
                        continue;
                    }

                    BlockPos chestPos = mutable.immutable();
                    Container container = NpcGoalHelper.getChestContainer(
                        level,
                        chestPos
                    );
                    if (container == null) continue;

                    // Index all items in this chest
                    Set<Item> seen = new HashSet<>();
                    for (int i = 0; i < container.getContainerSize(); i++) {
                        ItemStack stack = container.getItem(i);
                        if (!stack.isEmpty()) {
                            Item item = stack.getItem();
                            if (seen.add(item)) {
                                chestContentMap
                                    .computeIfAbsent(item, k ->
                                        new ArrayList<>()
                                    )
                                    .add(chestPos);
                            }
                        }
                    }
                }
            }
        }
    }

    /**
     * Find the best destination chest for an item: the one with the most
     * of the same item already (best consolidation).
     */
    @Nullable
    private BlockPos findBestDestination(List<BlockPos> candidates, Item item) {
        Level level = npc.level();
        BlockPos best = null;
        int bestCount = -1;

        for (BlockPos pos : candidates) {
            Container container = NpcGoalHelper.getChestContainer(level, pos);
            if (container == null) continue;

            // Check if chest has space and count matching items
            int count = 0;
            boolean hasSpace = false;
            for (int i = 0; i < container.getContainerSize(); i++) {
                ItemStack stack = container.getItem(i);
                if (stack.isEmpty()) {
                    hasSpace = true;
                } else if (stack.getItem() == item) {
                    count += stack.getCount();
                    if (stack.getCount() < stack.getMaxStackSize()) {
                        hasSpace = true;
                    }
                }
            }

            if (hasSpace && count > bestCount) {
                bestCount = count;
                best = pos;
            }
        }

        return best;
    }

    // ========== Utility ==========

    private void returnItemToSource() {
        if (
            heldItem == null || heldItem.isEmpty() || sourceChest == null
        ) return;

        Container container = NpcGoalHelper.getChestContainer(
            npc.level(),
            sourceChest
        );
        if (container != null) {
            NpcGoalHelper.insertIntoContainer(container, heldItem);
        }
        if (heldItem != null && !heldItem.isEmpty()) {
            npc.spawnAtLocation(heldItem);
        }
        this.heldItem = null;
    }

    private void awardXP(int amount) {
        if (this.master != null) {
        }
    }
}
