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

import com.tiedup.remake.entities.EntityLaborGuard;
import com.tiedup.remake.items.ItemTaser;
import com.tiedup.remake.items.ModItems;
import java.util.EnumSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import net.minecraft.server.level.ServerPlayer;
import net.minecraft.world.InteractionHand;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.entity.ai.attributes.Attributes;
import net.minecraft.world.entity.ai.goal.Goal;
import net.minecraft.world.entity.monster.*;
import net.minecraft.world.item.Item;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.item.Items;
import net.minecraft.world.phys.AABB;

/**
 * AI Goal for guard to hunt and kill monsters near the prisoner.
 *
 * Adapted from KidnapperHuntMonstersGoal for EntityDamsel-based guard.
 * Scans 16 blocks around the PRISONER (not the guard), deals 2x damage,
 * applies taser stun, and cancels creeper explosions.
 */
public class GuardHuntMonstersGoal extends Goal {

    private final EntityLaborGuard guard;
    private LivingEntity targetMonster;

    private static final int SCAN_RADIUS = 12;
    private static final int ATTACK_COOLDOWN = 15;
    private static final float MONSTER_DAMAGE_MULTIPLIER = 2.0f;

    /** Maps combat task drop items to the mob classes that drop them. */
    private static final Map<Item, Set<Class<? extends Monster>>> DROP_TO_MOB =
        Map.of(
            Items.ROTTEN_FLESH,
            Set.of(
                Zombie.class,
                Husk.class,
                Drowned.class,
                ZombieVillager.class
            ),
            Items.BONE,
            Set.of(Skeleton.class, WitherSkeleton.class, Stray.class),
            Items.STRING,
            Set.of(Spider.class, CaveSpider.class),
            Items.GUNPOWDER,
            Set.of(Creeper.class),
            Items.SPIDER_EYE,
            Set.of(Spider.class, CaveSpider.class, Witch.class),
            Items.ARROW,
            Set.of(Skeleton.class, Stray.class)
        );

    private int attackCooldown = 0;

    public GuardHuntMonstersGoal(EntityLaborGuard guard) {
        this.guard = guard;
        this.setFlags(EnumSet.of(Goal.Flag.MOVE, Goal.Flag.LOOK));
    }

    @Override
    public boolean canUse() {
        if (guard.isTiedUp()) {
            return false;
        }

        // Don't hunt if fighting a player target
        if (guard.getTarget() != null) {
            return false;
        }

        // Must have a prisoner to protect
        ServerPlayer prisoner = guard.getPrisoner();
        if (prisoner == null || !prisoner.isAlive()) {
            return false;
        }

        this.targetMonster = findNearestMonster(prisoner);
        return this.targetMonster != null;
    }

    @Override
    public boolean canContinueToUse() {
        if (this.targetMonster == null || !this.targetMonster.isAlive()) {
            return false;
        }

        if (guard.getTarget() != null) {
            return false;
        }

        if (guard.isTiedUp()) {
            return false;
        }

        // Abort if guard is getting too far from prisoner
        ServerPlayer prisoner = guard.getPrisoner();
        if (
            prisoner != null &&
            guard.distanceToSqr(prisoner) > SCAN_RADIUS * SCAN_RADIUS
        ) {
            return false;
        }

        // Abort if monster is no longer threatening
        if (this.targetMonster instanceof Monster monster && prisoner != null) {
            if (!isThreateningTarget(monster, prisoner)) {
                return false;
            }
        }

        return (
            guard.distanceToSqr(this.targetMonster) <
            SCAN_RADIUS * SCAN_RADIUS * 4
        );
    }

    @Override
    public void start() {
        this.attackCooldown = 0;

        guard.setItemInHand(
            InteractionHand.MAIN_HAND,
            new ItemStack(ModItems.TASER.get())
        );
    }

    @Override
    public void stop() {
        this.targetMonster = null;
        guard.getNavigation().stop();

        guard.setItemInHand(InteractionHand.MAIN_HAND, ItemStack.EMPTY);
    }

    @Override
    public void tick() {
        if (this.targetMonster == null) {
            return;
        }

        guard.getLookControl().setLookAt(this.targetMonster);

        double distSq = guard.distanceToSqr(this.targetMonster);

        if (distSq > 4.0) {
            guard.getNavigation().moveTo(this.targetMonster, 1.2);
        } else {
            guard.getNavigation().stop();

            if (this.attackCooldown <= 0) {
                attackMonsterWithBonus(this.targetMonster);
                this.attackCooldown = ATTACK_COOLDOWN;
            }
        }

        if (this.attackCooldown > 0) {
            this.attackCooldown--;
        }
    }

    private void attackMonsterWithBonus(LivingEntity target) {
        guard.swing(InteractionHand.MAIN_HAND);

        float baseDamage = (float) guard.getAttributeValue(
            Attributes.ATTACK_DAMAGE
        );
        float bonusDamage = baseDamage * MONSTER_DAMAGE_MULTIPLIER;

        // Attribute kill to prisoner so mobs drop player-kill loot (skulls, etc.)
        ServerPlayer prisoner = guard.getPrisoner();
        boolean damaged = target.hurt(
            prisoner != null
                ? guard.damageSources().playerAttack(prisoner)
                : guard.damageSources().mobAttack(guard),
            bonusDamage
        );

        if (damaged) {
            ItemStack heldItem = guard.getItemInHand(InteractionHand.MAIN_HAND);
            if (heldItem.getItem() instanceof ItemTaser taserItem) {
                taserItem.hurtEnemy(heldItem, target, guard);
            }

            if (target instanceof Creeper creeper) {
                creeper.setSwellDir(-1);
            }
        }
    }

    /**
     * Find nearest threatening monster around the prisoner.
     * Only targets monsters that are actively targeting the prisoner or the guard,
     * that are reachable (not underwater when we're on land),
     * and that are NOT task-relevant mobs during a combat task.
     */
    private LivingEntity findNearestMonster(ServerPlayer prisoner) {
        AABB searchBox = prisoner.getBoundingBox().inflate(SCAN_RADIUS);
        Set<Class<? extends Monster>> taskMobs = getTaskRelevantMobs(prisoner);

        List<Monster> monsters = guard
            .level()
            .getEntitiesOfClass(
                Monster.class,
                searchBox,
                m ->
                    m.isAlive() &&
                    !m.isSpectator() &&
                    !taskMobs.contains(m.getClass()) &&
                    isThreateningTarget(m, prisoner)
            );

        if (monsters.isEmpty()) {
            return null;
        }

        Monster nearest = null;
        double nearestDist = Double.MAX_VALUE;

        for (Monster m : monsters) {
            double dist = guard.distanceToSqr(m);
            if (dist < nearestDist) {
                nearestDist = dist;
                nearest = m;
            }
        }

        return nearest;
    }

    /**
     * Check if a monster is actually threatening the prisoner or the guard.
     */
    private boolean isThreateningTarget(
        Monster monster,
        ServerPlayer prisoner
    ) {
        // Must be targeting us or the prisoner
        LivingEntity monsterTarget = monster.getTarget();
        if (monsterTarget == null) {
            return false;
        }
        boolean targetingPrisoner = monsterTarget
            .getUUID()
            .equals(prisoner.getUUID());
        boolean targetingGuard = monsterTarget
            .getUUID()
            .equals(guard.getUUID());
        if (!targetingPrisoner && !targetingGuard) {
            return false;
        }

        // Skip if monster is in water and we're not (unreachable chase)
        if (monster.isInWater() && !guard.isInWater()) {
            return false;
        }

        return true;
    }

    /**
     * Get the set of mob classes that the prisoner's current task needs.
     * Returns empty set if no combat task is active.
     */
    private Set<Class<? extends Monster>> getTaskRelevantMobs(
        ServerPlayer prisoner
    ) {
        if (
            !(guard.level() instanceof
                    net.minecraft.server.level.ServerLevel serverLevel)
        ) {
            return Set.of();
        }
        com.tiedup.remake.prison.PrisonerManager manager =
            com.tiedup.remake.prison.PrisonerManager.get(serverLevel);
        com.tiedup.remake.prison.LaborRecord labor = manager.getLaborRecord(
            prisoner.getUUID()
        );
        if (
            labor.getPhase() !=
            com.tiedup.remake.prison.LaborRecord.WorkPhase.WORKING
        ) {
            return Set.of();
        }
        com.tiedup.remake.labor.LaborTask task = labor.getTask();
        if (task == null || !task.isCombatTask()) {
            return Set.of();
        }
        Set<Class<? extends Monster>> mobs = DROP_TO_MOB.get(
            task.getTargetItem()
        );
        return mobs != null ? mobs : Set.of();
    }
}
