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

import com.tiedup.remake.core.TiedUpMod;
import com.tiedup.remake.dialogue.DialogueBridge;
import com.tiedup.remake.entities.EntityMaster;
import com.tiedup.remake.v2.bondage.CollarHelper;
import com.tiedup.remake.v2.bondage.datadriven.DataDrivenBondageItem;
import net.minecraft.resources.ResourceLocation;
import com.tiedup.remake.state.PlayerBindState;
import com.tiedup.remake.v2.BodyRegionV2;
import com.tiedup.remake.v2.bondage.IV2BondageEquipment;
import com.tiedup.remake.v2.bondage.capability.V2EquipmentHelper;
import java.util.ArrayList;
import java.util.EnumSet;
import java.util.List;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.server.level.ServerPlayer;
import net.minecraft.sounds.SoundEvents;
import net.minecraft.sounds.SoundSource;
import net.minecraft.world.entity.ai.goal.Goal;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.phys.Vec3;

/**
 * AI Goal for EntityMaster to punish their pet player.
 *
 * Triggered when:
 * - Struggle attempt detected while master was watching
 * - Contraband found during inspection
 * - Task failure
 * - Pet attacks the master (dual punishment: choke + physical restraint)
 *
 * Punishment types (selected from available options):
 * - Choke collar activation (requires choke collar)
 * - Temporary blindfold (requires empty blindfold slot)
 * - Temporary gag (requires empty gag slot)
 * - Temporary mittens (requires empty mittens slot)
 * - Tighten restraints / apply armbinder (requires pet unbound)
 * - Cold shoulder (always available)
 * - Leash tug (always available)
 */
public class MasterPunishGoal extends Goal {

    private final EntityMaster master;

    /** Approach distance for punishment */
    private static final double PUNISH_DISTANCE = 2.0;

    /** Duration of punishment sequence (ticks) */
    private static final int PUNISH_DURATION = 80;

    /** Maximum choke duration before deactivating (ticks) - 3 seconds */
    private static final int MAX_CHOKE_TIME = 60;

    /** Duration of temporary punishment items (2 minutes) */
    private static final int TEMP_ITEM_DURATION = 2400;

    /** Fallback damage if no punishment method available */
    private static final float FALLBACK_DAMAGE = 2.0f;

    /** Delay between primary choke and secondary restraint (ticks) - 0.75s */
    private static final int SECONDARY_PUNISHMENT_DELAY = 15;

    /** Maximum punishment duration safety cap (ticks) - 8 seconds */
    private static final int MAX_PUNISH_DURATION = 160;

    /** NBT tags for temporary punishment items */
    private static final String NBT_TEMP_MASTER_EVENT = "tempMasterEvent";
    private static final String NBT_EXPIRATION_TIME = "expirationTime";

    private int punishTimer = 0;
    private boolean hasAppliedPunishment = false;
    private PunishmentType selectedPunishment = null;

    /** Timer for tracking choke duration - 0 means not choking */
    private int chokeActiveTimer = 0;

    /** Reference to the active choke collar for deactivation */
    private ItemStack activeChokeCollar = ItemStack.EMPTY;

    /** Leash tug timer */
    private int leashTugTimer = 0;

    // --- Attack punishment (dual: choke + physical restraint) ---

    /** Whether this is an attack-triggered dual punishment */
    private boolean isAttackPunishment = false;

    /** The physical restraint to apply after the choke */
    private PunishmentType secondaryPunishment = null;

    /** Whether the secondary restraint has been applied */
    private boolean hasAppliedSecondary = false;

    /** Tick count when the primary punishment was applied (for delay timing) */
    private int primaryAppliedAtTick = -1;

    public MasterPunishGoal(EntityMaster master) {
        this.master = master;
        this.setFlags(EnumSet.of(Goal.Flag.MOVE, Goal.Flag.LOOK));
    }

    @Override
    public boolean canUse() {
        return (
            master.getStateManager().getCurrentState() == MasterState.PUNISH &&
            master.hasPet() &&
            master.getPetPlayer() != null
        );
    }

    @Override
    public boolean canContinueToUse() {
        if (master.getStateManager().getCurrentState() != MasterState.PUNISH) {
            return false;
        }

        // Extend duration if waiting for secondary punishment, with safety cap
        if (
            isAttackPunishment &&
            !hasAppliedSecondary &&
            punishTimer < MAX_PUNISH_DURATION
        ) {
            return true;
        }

        return punishTimer < PUNISH_DURATION;
    }

    @Override
    public void start() {
        this.punishTimer = 0;
        this.hasAppliedPunishment = false;
        this.chokeActiveTimer = 0;
        this.activeChokeCollar = ItemStack.EMPTY;
        this.leashTugTimer = 0;
        this.isAttackPunishment = false;
        this.secondaryPunishment = null;
        this.hasAppliedSecondary = false;
        this.primaryAppliedAtTick = -1;

        // Check if this is an attack-triggered punishment (dual: choke + restraint)
        if (master.consumeAttackPunishment()) {
            this.isAttackPunishment = true;
            this.selectedPunishment = PunishmentType.CHOKE_COLLAR;
            this.secondaryPunishment = selectPhysicalRestraint();

            TiedUpMod.LOGGER.debug(
                "[MasterPunishGoal] {} starting ATTACK punishment: {} + {}",
                master.getNpcName(),
                selectedPunishment,
                secondaryPunishment
            );
        } else {
            this.selectedPunishment = selectPunishment();

            TiedUpMod.LOGGER.debug(
                "[MasterPunishGoal] {} starting punishment: {}",
                master.getNpcName(),
                selectedPunishment
            );
        }

        master.markEngagement();
    }

    @Override
    public void stop() {
        deactivateChoke();

        this.punishTimer = 0;
        this.hasAppliedPunishment = false;
        this.chokeActiveTimer = 0;
        this.activeChokeCollar = ItemStack.EMPTY;
        this.leashTugTimer = 0;
        this.selectedPunishment = null;
        this.isAttackPunishment = false;
        this.secondaryPunishment = null;
        this.hasAppliedSecondary = false;
        this.primaryAppliedAtTick = -1;

        if (master.hasPet()) {
            master.setMasterState(MasterState.FOLLOWING);
        }

        TiedUpMod.LOGGER.debug(
            "[MasterPunishGoal] {} punishment complete",
            master.getNpcName()
        );
    }

    @Override
    public void tick() {
        ServerPlayer pet = master.getPetPlayer();
        if (pet == null) return;

        master.getLookControl().setLookAt(pet, 30.0F, 30.0F);

        double distSq = master.distanceToSqr(pet);

        // Move close for punishment
        if (distSq > PUNISH_DISTANCE * PUNISH_DISTANCE) {
            master.getNavigation().moveTo(pet, 1.2);
        } else {
            master.getNavigation().stop();

            if (!hasAppliedPunishment) {
                applyPunishment(pet);
                hasAppliedPunishment = true;
                primaryAppliedAtTick = punishTimer;
            }
        }

        // Apply secondary punishment after delay (attack punishment only)
        if (
            isAttackPunishment &&
            !hasAppliedSecondary &&
            primaryAppliedAtTick >= 0 &&
            punishTimer >= primaryAppliedAtTick + SECONDARY_PUNISHMENT_DELAY
        ) {
            applySecondaryPunishment(pet);
            hasAppliedSecondary = true;
        }

        // Track choke duration
        if (chokeActiveTimer > 0) {
            chokeActiveTimer++;
            if (chokeActiveTimer >= MAX_CHOKE_TIME) {
                deactivateChoke();
            }
        }

        // Track leash tug
        if (leashTugTimer > 0) {
            leashTugTimer--;
            if (leashTugTimer <= 0) {
                master.detachLeashFromPet();
            }
        }

        punishTimer++;
    }

    /**
     * Select a punishment type from available options (normal punishments).
     */
    private PunishmentType selectPunishment() {
        // Check for forced punishment (e.g., from other systems)
        PunishmentType forced = master.consumeForcedPunishment();
        if (forced != null) return forced;

        ServerPlayer pet = master.getPetPlayer();
        if (pet == null) return PunishmentType.COLD_SHOULDER;

        PlayerBindState bindState = PlayerBindState.getInstance(pet);
        List<PunishmentType> available = new ArrayList<>();

        // CHOKE: only if pet has choke collar
        if (bindState != null && bindState.hasCollar()) {
            ItemStack collar = bindState.getEquipment(BodyRegionV2.NECK);
            if (CollarHelper.isChokeCollar(collar)) {
                available.add(PunishmentType.CHOKE_COLLAR);
            }
        }

        // BLINDFOLD: only if eyes region is empty
        if (!V2EquipmentHelper.isRegionOccupied(pet, BodyRegionV2.EYES)) {
            available.add(PunishmentType.BLINDFOLD);
        }

        // GAG: only if mouth region is empty
        if (!V2EquipmentHelper.isRegionOccupied(pet, BodyRegionV2.MOUTH)) {
            available.add(PunishmentType.GAG);
        }

        // MITTENS: only if hands region is empty
        if (!V2EquipmentHelper.isRegionOccupied(pet, BodyRegionV2.HANDS)) {
            available.add(PunishmentType.MITTENS);
        }

        // TIGHTEN: only if pet is not already bound
        if (!V2EquipmentHelper.isRegionOccupied(pet, BodyRegionV2.ARMS)) {
            available.add(PunishmentType.TIGHTEN_RESTRAINTS);
        }

        // COLD_SHOULDER and LEASH_TUG: always available
        available.add(PunishmentType.COLD_SHOULDER);
        available.add(PunishmentType.LEASH_TUG);

        return available.get(master.getRandom().nextInt(available.size()));
    }

    /**
     * Select a physical restraint for attack punishment secondary.
     * Excludes CHOKE_COLLAR (already primary) and COLD_SHOULDER (not physical).
     * Always includes LEASH_TUG as fallback when all equipment slots are full.
     */
    private PunishmentType selectPhysicalRestraint() {
        ServerPlayer pet = master.getPetPlayer();
        if (pet == null) return PunishmentType.LEASH_TUG;

        List<PunishmentType> available = new ArrayList<>();

        // BLINDFOLD: only if eyes region is empty
        if (!V2EquipmentHelper.isRegionOccupied(pet, BodyRegionV2.EYES)) {
            available.add(PunishmentType.BLINDFOLD);
        }

        // GAG: only if mouth region is empty
        if (!V2EquipmentHelper.isRegionOccupied(pet, BodyRegionV2.MOUTH)) {
            available.add(PunishmentType.GAG);
        }

        // MITTENS: only if hands region is empty
        if (!V2EquipmentHelper.isRegionOccupied(pet, BodyRegionV2.HANDS)) {
            available.add(PunishmentType.MITTENS);
        }

        // TIGHTEN: only if pet is not already bound
        if (!V2EquipmentHelper.isRegionOccupied(pet, BodyRegionV2.ARMS)) {
            available.add(PunishmentType.TIGHTEN_RESTRAINTS);
        }

        // LEASH_TUG: always available as fallback
        available.add(PunishmentType.LEASH_TUG);

        return available.get(master.getRandom().nextInt(available.size()));
    }

    /**
     * Apply the selected punishment to the pet.
     */
    private void applyPunishment(ServerPlayer pet) {
        if (selectedPunishment == null) {
            selectedPunishment = PunishmentType.COLD_SHOULDER;
        }

        // Send dialogue
        DialogueBridge.talkTo(master, pet, selectedPunishment.getDialogueId());

        switch (selectedPunishment) {
            case CHOKE_COLLAR -> applyChoke(pet);
            case BLINDFOLD -> applyTempAccessory(pet, BodyRegionV2.EYES);
            case GAG -> applyTempAccessory(pet, BodyRegionV2.MOUTH);
            case MITTENS -> applyTempAccessory(pet, BodyRegionV2.HANDS);
            case TIGHTEN_RESTRAINTS -> applyTighten(pet);
            case COLD_SHOULDER -> applyColdShoulder();
            case LEASH_TUG -> applyLeashTug(pet);
        }

        TiedUpMod.LOGGER.info(
            "[MasterPunishGoal] {} applied {} to {}",
            master.getNpcName(),
            selectedPunishment,
            pet.getName().getString()
        );
    }

    /**
     * Apply the secondary physical restraint after the choke (attack punishment only).
     */
    private void applySecondaryPunishment(ServerPlayer pet) {
        if (secondaryPunishment == null) return;

        // Send dialogue for the secondary punishment
        DialogueBridge.talkTo(master, pet, secondaryPunishment.getDialogueId());

        switch (secondaryPunishment) {
            case BLINDFOLD -> applyTempAccessory(pet, BodyRegionV2.EYES);
            case GAG -> applyTempAccessory(pet, BodyRegionV2.MOUTH);
            case MITTENS -> applyTempAccessory(pet, BodyRegionV2.HANDS);
            case TIGHTEN_RESTRAINTS -> applyTighten(pet);
            case LEASH_TUG -> applyLeashTug(pet);
            default -> {
            } // CHOKE_COLLAR and COLD_SHOULDER excluded from secondary
        }

        TiedUpMod.LOGGER.info(
            "[MasterPunishGoal] {} applied secondary {} to {}",
            master.getNpcName(),
            secondaryPunishment,
            pet.getName().getString()
        );
    }

    /**
     * Apply choke collar punishment (existing logic).
     */
    private void applyChoke(ServerPlayer pet) {
        PlayerBindState bindState = PlayerBindState.getInstance(pet);
        if (bindState != null && bindState.hasCollar()) {
            ItemStack collar = bindState.getEquipment(BodyRegionV2.NECK);
            if (CollarHelper.isChokeCollar(collar)) {
                CollarHelper.setChoking(collar, true);
                this.activeChokeCollar = collar;
                this.chokeActiveTimer = 1;

                pet
                    .level()
                    .playSound(
                        null,
                        pet.getX(),
                        pet.getY(),
                        pet.getZ(),
                        SoundEvents.PLAYER_HURT,
                        SoundSource.HOSTILE,
                        0.8f,
                        0.5f + master.getRandom().nextFloat() * 0.2f
                    );
                return;
            }
        }
        // Fallback
        pet.hurt(pet.damageSources().magic(), FALLBACK_DAMAGE);
    }

    /**
     * Apply a temporary accessory (blindfold, gag, or mittens) as punishment.
     * Uses the same temp-item pattern as MasterRandomEventGoal.
     */
    private void applyTempAccessory(ServerPlayer pet, BodyRegionV2 region) {
        ItemStack accessory = createAccessory(region);
        if (accessory.isEmpty()) return;

        // Mark as temporary with expiration
        long expirationTime = master.level().getGameTime() + TEMP_ITEM_DURATION;
        CompoundTag tag = accessory.getOrCreateTag();
        tag.putBoolean(NBT_TEMP_MASTER_EVENT, true);
        tag.putLong(NBT_EXPIRATION_TIME, expirationTime);
        tag.putUUID("masterUUID", master.getUUID());

        IV2BondageEquipment equip = V2EquipmentHelper.getEquipment(pet);
        if (equip != null) {
            equip.setInRegion(region, accessory);
            V2EquipmentHelper.sync(pet);
        }

        pet
            .level()
            .playSound(
                null,
                pet.getX(),
                pet.getY(),
                pet.getZ(),
                SoundEvents.ARMOR_EQUIP_LEATHER,
                SoundSource.HOSTILE,
                1.0f,
                0.8f
            );
    }

    /**
     * Create an accessory item for the given body region.
     */
    private ItemStack createAccessory(BodyRegionV2 region) {
        return switch (region) {
            case EYES -> DataDrivenBondageItem.createStack(
                ResourceLocation.fromNamespaceAndPath("tiedup", "classic_blindfold")
            );
            case MOUTH -> DataDrivenBondageItem.createStack(
                ResourceLocation.fromNamespaceAndPath("tiedup", "ball_gag")
            );
            case HANDS -> DataDrivenBondageItem.createStack(
                ResourceLocation.fromNamespaceAndPath("tiedup", "leather_mittens")
            );
            default -> ItemStack.EMPTY;
        };
    }

    /**
     * Apply armbinder as punishment.
     */
    private void applyTighten(ServerPlayer pet) {
        ItemStack armbinder = DataDrivenBondageItem.createStack(
            ResourceLocation.fromNamespaceAndPath("tiedup", "armbinder")
        );

        // Mark as temporary
        long expirationTime = master.level().getGameTime() + TEMP_ITEM_DURATION;
        CompoundTag tag = armbinder.getOrCreateTag();
        tag.putBoolean(NBT_TEMP_MASTER_EVENT, true);
        tag.putLong(NBT_EXPIRATION_TIME, expirationTime);
        tag.putUUID("masterUUID", master.getUUID());

        IV2BondageEquipment equip = V2EquipmentHelper.getEquipment(pet);
        if (equip != null) {
            equip.setInRegion(BodyRegionV2.ARMS, armbinder);
            V2EquipmentHelper.sync(pet);
        }

        pet
            .level()
            .playSound(
                null,
                pet.getX(),
                pet.getY(),
                pet.getZ(),
                SoundEvents.ARMOR_EQUIP_CHAIN,
                SoundSource.HOSTILE,
                1.0f,
                0.7f
            );
    }

    /**
     * Apply cold shoulder - master ignores pet interactions.
     */
    private void applyColdShoulder() {
        master.startColdShoulder(
            PunishmentType.COLD_SHOULDER.getDurationTicks()
        );
    }

    /**
     * Apply leash tug - yank pet toward master.
     */
    private void applyLeashTug(ServerPlayer pet) {
        // Attach temp leash
        master.attachLeashToPet();
        this.leashTugTimer = PunishmentType.LEASH_TUG.getDurationTicks();

        // Apply velocity toward master
        Vec3 direction = master.position().subtract(pet.position()).normalize();
        pet.setDeltaMovement(direction.scale(1.2));
        pet.hurtMarked = true;

        pet
            .level()
            .playSound(
                null,
                pet.getX(),
                pet.getY(),
                pet.getZ(),
                SoundEvents.CHAIN_HIT,
                SoundSource.HOSTILE,
                1.0f,
                1.0f
            );
    }

    /**
     * Deactivate the choke collar if active.
     */
    private void deactivateChoke() {
        if (
            !activeChokeCollar.isEmpty() &&
            CollarHelper.isChokeCollar(activeChokeCollar)
        ) {
            CollarHelper.setChoking(activeChokeCollar, false);
        }
        this.activeChokeCollar = ItemStack.EMPTY;
        this.chokeActiveTimer = 0;
    }
}
