package com.tiedup.remake.compat.mca.ai;

import com.tiedup.remake.compat.mca.MCABondageManager;
import com.tiedup.remake.compat.mca.MCACompat;
import com.tiedup.remake.compat.mca.ai.goals.MCAFleeGoal;
import com.tiedup.remake.compat.mca.ai.goals.MCAFollowCaptorGoal;
import com.tiedup.remake.compat.mca.ai.goals.MCAPanicGoal;
import com.tiedup.remake.compat.mca.ai.goals.MCAStayGoal;
import com.tiedup.remake.compat.mca.personality.MCAPersonality;
import com.tiedup.remake.compat.mca.personality.MCAPersonalityManager;
import com.tiedup.remake.core.TiedUpMod;
import com.tiedup.remake.state.IBondageState;
import com.tiedup.remake.util.RestraintEffectUtils;
import java.lang.ref.WeakReference;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.entity.Mob;
import net.minecraft.world.entity.PathfinderMob;
import net.minecraft.world.entity.ai.goal.FloatGoal;
import net.minecraft.world.entity.ai.goal.Goal;
import net.minecraft.world.entity.ai.goal.WrappedGoal;
import net.minecraft.world.entity.player.Player;

/**
 * Controls AI behavior for a single MCA villager during bondage.
 *
 * Non-invasive design:
 * - Stores original AI state before modification
 * - Restores fully when freed
 * - Configurable per capture state
 *
 * The controller manages transitions between AI levels,
 * injecting and removing goals as needed.
 */
public class MCABondageAIController {

    private final WeakReference<LivingEntity> villagerRef;
    private MCABondageAILevel currentLevel = MCABondageAILevel.NONE;

    // Backup of original goals (for restoration in OVERRIDE mode)
    private Set<Goal> originalGoals = null;
    private boolean mcaBrainSuspended = false;

    // TiedUp-specific goals we've added
    private final List<Goal> injectedGoals = new ArrayList<>();

    // Personality for behavior modification
    private MCAPersonality personality = MCAPersonality.UNKNOWN;

    /**
     * Create AI controller for an MCA villager.
     *
     * @param villager The MCA villager entity
     */
    public MCABondageAIController(LivingEntity villager) {
        this.villagerRef = new WeakReference<>(villager);
    }

    /**
     * Get the current AI control level.
     */
    public MCABondageAILevel getCurrentLevel() {
        return currentLevel;
    }

    /**
     * Get the villager this controller manages.
     */
    public LivingEntity getVillager() {
        return villagerRef.get();
    }

    /**
     * Get the villager's personality.
     */
    public MCAPersonality getPersonality() {
        return personality;
    }

    /**
     * Set the villager's personality.
     * This affects AI behavior (flee speed, panic duration, etc.)
     */
    public void setPersonality(MCAPersonality personality) {
        this.personality =
            personality != null ? personality : MCAPersonality.UNKNOWN;
    }

    // LEVEL MANAGEMENT

    /**
     * Recalculate and apply appropriate AI level based on current state.
     * Call this after any bondage state change.
     */
    public void updateAILevel() {
        LivingEntity villager = villagerRef.get();
        if (villager == null) return;

        IBondageState state = MCABondageManager.getInstance().getKidnappedState(
            villager
        );
        if (state == null) {
            setLevel(MCABondageAILevel.NONE);
            return;
        }

        // Determine level from state
        MCABondageAILevel newLevel = calculateLevel(state);
        setLevel(newLevel);
    }

    /**
     * Calculate the appropriate AI level based on bondage state.
     */
    private MCABondageAILevel calculateLevel(IBondageState state) {
        // OVERRIDE: Fully tied (arms AND legs) OR gagged AND blindfolded
        if (state.hasArmsBound() && state.hasLegsBound()) {
            return MCABondageAILevel.OVERRIDE;
        }
        if (state.isGagged() && state.isBlindfolded()) {
            return MCABondageAILevel.OVERRIDE;
        }

        // MODIFIED: Tied (arms or legs)
        if (state.isTiedUp()) {
            return MCABondageAILevel.MODIFIED;
        }

        // BASIC: Collar only
        if (state.hasCollar()) {
            return MCABondageAILevel.BASIC;
        }

        // NONE: No restrictions
        return MCABondageAILevel.NONE;
    }

    /**
     * Set a specific AI level.
     * Handles transitions between levels properly.
     *
     * @param level The level to set
     */
    public void setLevel(MCABondageAILevel level) {
        if (level == currentLevel) return;

        LivingEntity villager = villagerRef.get();
        if (villager == null) return;

        TiedUpMod.LOGGER.debug(
            "[MCA AI] Level change for {}: {} -> {}",
            villager.getName().getString(),
            currentLevel,
            level
        );

        // Clean up current level effects
        transitionFromLevel(currentLevel);

        // Apply new level effects
        MCABondageAILevel oldLevel = currentLevel;
        currentLevel = level;
        transitionToLevel(level, oldLevel);
    }

    /**
     * Clean up effects from the previous level.
     */
    private void transitionFromLevel(MCABondageAILevel level) {
        switch (level) {
            case OVERRIDE -> {
                restoreMCABrain();
                restoreOriginalGoals();
            }
            case MODIFIED -> removeInjectedGoals();
            case BASIC -> removeSpeedReduction();
            case NONE -> {
                /* nothing to clean */
            }
        }
    }

    /**
     * Apply effects for the new level.
     */
    private void transitionToLevel(
        MCABondageAILevel level,
        MCABondageAILevel oldLevel
    ) {
        switch (level) {
            case NONE -> {
                /* nothing to apply */
            }
            case BASIC -> applySpeedReduction();
            case MODIFIED -> {
                if (!oldLevel.hasSpeedReduction()) {
                    applySpeedReduction();
                }
                injectBondageGoals();
            }
            case OVERRIDE -> {
                if (!oldLevel.hasSpeedReduction()) {
                    applySpeedReduction();
                }
                backupAndClearGoals();
                suspendMCABrain();
                injectOverrideGoals();
            }
        }
    }

    // SPEED REDUCTION

    private void applySpeedReduction() {
        LivingEntity villager = villagerRef.get();
        if (villager == null || villager.level().isClientSide()) return;

        RestraintEffectUtils.applyBindSpeedReduction(villager);

        // Stop current navigation
        if (villager instanceof Mob mob) {
            mob.getNavigation().stop();
        }

        TiedUpMod.LOGGER.debug(
            "[MCA AI] Applied speed reduction to {}",
            villager.getName().getString()
        );
    }

    private void removeSpeedReduction() {
        LivingEntity villager = villagerRef.get();
        if (villager == null || villager.level().isClientSide()) return;

        RestraintEffectUtils.removeBindSpeedReduction(villager);

        TiedUpMod.LOGGER.debug(
            "[MCA AI] Removed speed reduction from {}",
            villager.getName().getString()
        );
    }

    // GOAL INJECTION (MODIFIED LEVEL)

    private void injectBondageGoals() {
        LivingEntity villager = villagerRef.get();
        if (!(villager instanceof PathfinderMob mob)) return;

        IBondageState state = MCABondageManager.getInstance().getKidnappedState(
            villager
        );
        if (state == null) return;

        // Get combined flee/panic speed from personality
        float fleeMultiplier =
            MCAPersonalityManager.getInstance().getCombinedCompliance(villager);
        // Invert compliance for flee speed (compliant = slow flee, rebellious = fast flee)
        float fleeSpeedMod = 2.0f - fleeMultiplier; // Range ~0.5 to ~1.7
        fleeSpeedMod = Math.max(0.5f, Math.min(1.5f, fleeSpeedMod));

        // Priority 1: Follow captor when leashed
        MCAFollowCaptorGoal followGoal = new MCAFollowCaptorGoal(
            mob,
            1.0,
            10.0f,
            2.0f
        );
        mob.goalSelector.addGoal(1, followGoal);
        injectedGoals.add(followGoal);

        // Priority 2: Panic when being tied (personality affects speed)
        double panicSpeed = 1.2 * personality.getFleeSpeedMultiplier();
        MCAPanicGoal panicGoal = new MCAPanicGoal(mob, panicSpeed);
        mob.goalSelector.addGoal(2, panicGoal);
        injectedGoals.add(panicGoal);

        // Priority 3: Flee from players if not collared (personality affects speed)
        if (!state.hasCollar()) {
            double walkSpeed = 1.0 * personality.getFleeSpeedMultiplier();
            double sprintSpeed = 1.2 * personality.getFleeSpeedMultiplier();
            MCAFleeGoal fleeGoal = new MCAFleeGoal(
                mob,
                Player.class,
                8.0f,
                walkSpeed,
                sprintSpeed
            );
            mob.goalSelector.addGoal(3, fleeGoal);
            injectedGoals.add(fleeGoal);
        }

        TiedUpMod.LOGGER.debug(
            "[MCA AI] Injected {} bondage goals for {} (personality: {}, flee mod: {})",
            injectedGoals.size(),
            villager.getName().getString(),
            personality.getMcaId(),
            String.format("%.2f", personality.getFleeSpeedMultiplier())
        );
    }

    private void removeInjectedGoals() {
        LivingEntity villager = villagerRef.get();
        if (!(villager instanceof Mob mob)) return;

        for (Goal goal : injectedGoals) {
            mob.goalSelector.removeGoal(goal);
        }

        int count = injectedGoals.size();
        injectedGoals.clear();

        TiedUpMod.LOGGER.debug(
            "[MCA AI] Removed {} injected goals from {}",
            count,
            villager.getName().getString()
        );
    }

    // FULL AI OVERRIDE

    private void backupAndClearGoals() {
        LivingEntity villager = villagerRef.get();
        if (!(villager instanceof Mob mob)) return;

        // Backup original goals
        originalGoals = new HashSet<>();
        for (WrappedGoal wrappedGoal : mob.goalSelector.getAvailableGoals()) {
            originalGoals.add(wrappedGoal.getGoal());
        }

        // Clear all goals
        mob.goalSelector.removeAllGoals(g -> true);

        TiedUpMod.LOGGER.debug(
            "[MCA AI] Backed up {} goals for {}",
            originalGoals.size(),
            villager.getName().getString()
        );
    }

    private void restoreOriginalGoals() {
        if (originalGoals == null) return;

        LivingEntity villager = villagerRef.get();
        if (!(villager instanceof Mob mob)) return;

        // Remove our override goals first
        for (Goal goal : injectedGoals) {
            mob.goalSelector.removeGoal(goal);
        }
        injectedGoals.clear();

        // Restore original goals
        // Note: We don't know original priorities, so use default of 5
        for (Goal goal : originalGoals) {
            mob.goalSelector.addGoal(5, goal);
        }

        int count = originalGoals.size();
        originalGoals = null;

        TiedUpMod.LOGGER.debug(
            "[MCA AI] Restored {} original goals for {}",
            count,
            villager.getName().getString()
        );
    }

    private void injectOverrideGoals() {
        LivingEntity villager = villagerRef.get();
        if (!(villager instanceof Mob mob)) return;

        // Essential goals only
        // Priority 0: Always float (don't drown)
        FloatGoal floatGoal = new FloatGoal(mob);
        mob.goalSelector.addGoal(0, floatGoal);
        injectedGoals.add(floatGoal);

        // Priority 1: Stay in place
        MCAStayGoal stayGoal = new MCAStayGoal(mob);
        mob.goalSelector.addGoal(1, stayGoal);
        injectedGoals.add(stayGoal);

        TiedUpMod.LOGGER.debug(
            "[MCA AI] Injected override goals for {}",
            villager.getName().getString()
        );
    }

    // MCA BRAIN CONTROL

    /**
     * Suspend MCA's brain activities (for OVERRIDE level).
     * Uses reflection to access VillagerBrain if available.
     */
    private void suspendMCABrain() {
        LivingEntity villager = villagerRef.get();
        if (villager == null) return;

        try {
            // MCA villagers implement VillagerLike which has getVillagerBrain()
            // We try to set their move state to STAY
            Method getVillagerBrain = villager
                .getClass()
                .getMethod("getVillagerBrain");
            Object mcaBrain = getVillagerBrain.invoke(villager);

            if (mcaBrain != null) {
                // Try to find setMoveState method
                // MCA uses MoveState enum with values like MOVE, STAY, etc.
                for (Method method : mcaBrain.getClass().getMethods()) {
                    if (method.getName().equals("setMoveState")) {
                        // Find the MoveState enum
                        Class<?>[] paramTypes = method.getParameterTypes();
                        if (paramTypes.length >= 1 && paramTypes[0].isEnum()) {
                            Object[] enumConstants =
                                paramTypes[0].getEnumConstants();
                            // Find STAY constant
                            for (Object constant : enumConstants) {
                                if (constant.toString().equals("STAY")) {
                                    method.invoke(mcaBrain, constant, null);
                                    mcaBrainSuspended = true;
                                    TiedUpMod.LOGGER.debug(
                                        "[MCA AI] Suspended brain for {}",
                                        villager.getName().getString()
                                    );
                                    return;
                                }
                            }
                        }
                    }
                }
            }
        } catch (Exception e) {
            TiedUpMod.LOGGER.debug(
                "[MCA AI] Could not suspend brain for {}: {}",
                villager.getName().getString(),
                e.getMessage()
            );
        }
    }

    /**
     * Restore MCA's brain to normal state.
     */
    private void restoreMCABrain() {
        if (!mcaBrainSuspended) return;

        LivingEntity villager = villagerRef.get();
        if (villager == null) return;

        try {
            Method getVillagerBrain = villager
                .getClass()
                .getMethod("getVillagerBrain");
            Object mcaBrain = getVillagerBrain.invoke(villager);

            if (mcaBrain != null) {
                for (Method method : mcaBrain.getClass().getMethods()) {
                    if (method.getName().equals("setMoveState")) {
                        Class<?>[] paramTypes = method.getParameterTypes();
                        if (paramTypes.length >= 1 && paramTypes[0].isEnum()) {
                            Object[] enumConstants =
                                paramTypes[0].getEnumConstants();
                            // Find MOVE constant
                            for (Object constant : enumConstants) {
                                if (constant.toString().equals("MOVE")) {
                                    method.invoke(mcaBrain, constant, null);
                                    mcaBrainSuspended = false;
                                    TiedUpMod.LOGGER.debug(
                                        "[MCA AI] Restored brain for {}",
                                        villager.getName().getString()
                                    );
                                    return;
                                }
                            }
                        }
                    }
                }
            }
        } catch (Exception e) {
            TiedUpMod.LOGGER.debug(
                "[MCA AI] Could not restore brain for {}: {}",
                villager.getName().getString(),
                e.getMessage()
            );
        }

        mcaBrainSuspended = false;
    }

    // CLEANUP

    /**
     * Clean up all AI modifications.
     * Call when villager is freed or removed.
     */
    public void cleanup() {
        // Restore from current level
        transitionFromLevel(currentLevel);
        currentLevel = MCABondageAILevel.NONE;

        // Clear any remaining state
        injectedGoals.clear();
        originalGoals = null;
        mcaBrainSuspended = false;

        LivingEntity villager = villagerRef.get();
        if (villager != null) {
            TiedUpMod.LOGGER.debug(
                "[MCA AI] Cleaned up controller for {}",
                villager.getName().getString()
            );
        }
    }
}
