TiedUp-/src/main/java/com/tiedup/remake/entities/KidnapperItemSelector.java

package com.tiedup.remake.entities;

import com.tiedup.remake.items.base.*;
import com.tiedup.remake.v2.bondage.datadriven.DataDrivenBondageItem;
import java.util.Random;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.world.item.ItemStack;
import org.jetbrains.annotations.Nullable;

/**
 * Helper class for selecting themed items for kidnappers.
 * Handles probability-based item selection and color application.
 *
 * Item selection order (most to least common):
 * 1. Bind (arms) - 100% (always)
 * 2. Gag - 50%
 * 3. Mittens - 40%
 * 4. Earplugs - 30%
 * 5. Blindfold - 20% (last, most restrictive)
 *
 * Elite kidnappers get significant bonuses to all probabilities.
 */
public class KidnapperItemSelector {

    private static final Random RANDOM = new Random();

    /** NBT key for item color */
    public static final String NBT_ITEM_COLOR = "ItemColor";

    // === BASE PROBABILITIES (Regular Kidnapper) ===
    public static final double PROB_BIND = 1.0; // 100%
    public static final double PROB_GAG = 0.5; // 50%
    public static final double PROB_MITTENS = 0.4; // 40%
    public static final double PROB_EARPLUGS = 0.3; // 30%
    public static final double PROB_BLINDFOLD = 0.2; // 20%

    // === ELITE KIDNAPPER BONUSES ===
    public static final double ELITE_GAG_BONUS = 0.5; // +50% = 100%
    public static final double ELITE_MITTENS_BONUS = 0.4; // +40% = 80%
    public static final double ELITE_EARPLUGS_BONUS = 0.3; // +30% = 60%
    public static final double ELITE_BLINDFOLD_BONUS = 0.2; // +20% = 40%

    // === ARCHER KIDNAPPER PENALTIES (relies on arrows) ===
    public static final double ARCHER_GAG_PENALTY = 0.2; // -20% = 30%
    public static final double ARCHER_MITTENS_PENALTY = 0.2; // -20% = 20%
    public static final double ARCHER_EARPLUGS_PENALTY = 0.15; // -15% = 15%
    public static final double ARCHER_BLINDFOLD_PENALTY = 0.1; // -10% = 10%

    /**
     * Result of item selection for a kidnapper.
     * Contains theme, color, and all selected items.
     */
    public static class SelectionResult {

        public final KidnapperTheme theme;
        public final ItemColor color; // null if theme doesn't support colors
        public final ItemStack bind;
        public final ItemStack gag;
        public final ItemStack mittens;
        public final ItemStack earplugs;
        public final ItemStack blindfold;

        public SelectionResult(
            KidnapperTheme theme,
            ItemColor color,
            ItemStack bind,
            ItemStack gag,
            ItemStack mittens,
            ItemStack earplugs,
            ItemStack blindfold
        ) {
            this.theme = theme;
            this.color = color;
            this.bind = bind;
            this.gag = gag;
            this.mittens = mittens;
            this.earplugs = earplugs;
            this.blindfold = blindfold;
        }

        /**
         * Check if this selection has a gag.
         */
        public boolean hasGag() {
            return !gag.isEmpty();
        }

        /**
         * Check if this selection has mittens.
         */
        public boolean hasMittens() {
            return !mittens.isEmpty();
        }

        /**
         * Check if this selection has earplugs.
         */
        public boolean hasEarplugs() {
            return !earplugs.isEmpty();
        }

        /**
         * Check if this selection has a blindfold.
         */
        public boolean hasBlindfold() {
            return !blindfold.isEmpty();
        }
    }

    /**
     * Select items for a regular kidnapper.
     */
    public static SelectionResult selectForKidnapper() {
        return selectItems(false, false);
    }

    /**
     * Select items for an elite kidnapper (higher probabilities).
     */
    public static SelectionResult selectForEliteKidnapper() {
        return selectItems(true, false);
    }

    /**
     * Select items for an archer kidnapper (lower probabilities - relies on arrows).
     */
    public static SelectionResult selectForArcherKidnapper() {
        return selectItems(false, true);
    }

    /**
     * Calculate adjusted probability based on kidnapper type.
     *
     * @param baseProb     Base probability for the item
     * @param eliteBonus   Bonus probability for elite kidnappers
     * @param archerPenalty Penalty probability for archer kidnappers
     * @param isElite      Whether the kidnapper is elite
     * @param isArcher     Whether the kidnapper is an archer
     * @return Adjusted probability
     */
    private static double getAdjustedProbability(
        double baseProb,
        double eliteBonus,
        double archerPenalty,
        boolean isElite,
        boolean isArcher
    ) {
        double prob = baseProb;
        if (isElite) prob += eliteBonus;
        if (isArcher) prob -= archerPenalty;
        return prob;
    }

    /**
     * Internal item selection logic.
     */
    private static SelectionResult selectItems(
        boolean isElite,
        boolean isArcher
    ) {
        // 1. Select random theme
        KidnapperTheme theme = KidnapperTheme.getRandomWeighted();

        // 2. Select color (if theme supports it)
        // Filter out colors that don't have textures for this theme's bind
        ItemColor color = theme.supportsColor()
            ? getValidColorForBind(theme.getBind())
            : null;

        // 3. Create bind (always)
        ItemStack bind = createBind(theme.getBind(), color);

        // 4. Roll for gag (randomly selected from theme's compatible gags)
        ItemStack gag = ItemStack.EMPTY;
        double gagProb = getAdjustedProbability(
            PROB_GAG,
            ELITE_GAG_BONUS,
            ARCHER_GAG_PENALTY,
            isElite,
            isArcher
        );
        if (RANDOM.nextDouble() < gagProb) {
            gag = createGag(theme.getRandomGag(), color);
        }

        // 5. Roll for mittens (same for all themes)
        ItemStack mittens = ItemStack.EMPTY;
        double mittensProb = getAdjustedProbability(
            PROB_MITTENS,
            ELITE_MITTENS_BONUS,
            ARCHER_MITTENS_PENALTY,
            isElite,
            isArcher
        );
        if (RANDOM.nextDouble() < mittensProb) {
            mittens = createMittens();
        }

        // 6. Roll for earplugs (same for all themes)
        ItemStack earplugs = ItemStack.EMPTY;
        double earplugsProb = getAdjustedProbability(
            PROB_EARPLUGS,
            ELITE_EARPLUGS_BONUS,
            ARCHER_EARPLUGS_PENALTY,
            isElite,
            isArcher
        );
        if (RANDOM.nextDouble() < earplugsProb) {
            earplugs = createEarplugs();
        }

        // 7. Roll for blindfold (last, most restrictive - randomly selected)
        ItemStack blindfold = ItemStack.EMPTY;
        double blindfoldProb = getAdjustedProbability(
            PROB_BLINDFOLD,
            ELITE_BLINDFOLD_BONUS,
            ARCHER_BLINDFOLD_PENALTY,
            isElite,
            isArcher
        );
        if (theme.hasBlindfolds() && RANDOM.nextDouble() < blindfoldProb) {
            blindfold = createBlindfold(theme.getRandomBlindfold(), color);
        }

        return new SelectionResult(
            theme,
            color,
            bind,
            gag,
            mittens,
            earplugs,
            blindfold
        );
    }

    // ITEM CREATION METHODS

    /**
     * Create a bind ItemStack with optional color.
     */
    public static ItemStack createBind(
        BindVariant variant,
        @Nullable ItemColor color
    ) {
        ItemStack stack = DataDrivenBondageItem.createStack(
            new ResourceLocation("tiedup", variant.getRegistryName())
        );
        if (color != null && variant.supportsColor()) {
            applyColor(stack, color);
        }
        return stack;
    }

    /**
     * Create a gag ItemStack with optional color.
     * Validates that the color has a texture for this gag variant.
     */
    public static ItemStack createGag(
        GagVariant variant,
        @Nullable ItemColor color
    ) {
        ItemStack stack = DataDrivenBondageItem.createStack(
            new ResourceLocation("tiedup", variant.getRegistryName())
        );
        if (
            color != null &&
            variant.supportsColor() &&
            isColorValidForGag(color, variant)
        ) {
            applyColor(stack, color);
        }
        return stack;
    }

    /**
     * Create a blindfold ItemStack with optional color.
     * Validates that the color has a texture for this blindfold variant.
     */
    public static ItemStack createBlindfold(
        BlindfoldVariant variant,
        @Nullable ItemColor color
    ) {
        ItemStack stack = DataDrivenBondageItem.createStack(
            new ResourceLocation("tiedup", variant.getRegistryName())
        );
        if (
            color != null &&
            variant.supportsColor() &&
            isColorValidForBlindfold(color, variant)
        ) {
            applyColor(stack, color);
        }
        return stack;
    }

    /**
     * Create mittens ItemStack.
     * Mittens don't have color variants.
     */
    public static ItemStack createMittens() {
        return DataDrivenBondageItem.createStack(
            new ResourceLocation("tiedup", "leather_mittens")
        );
    }

    /**
     * Create earplugs ItemStack.
     * Earplugs don't have color variants.
     */
    public static ItemStack createEarplugs() {
        return DataDrivenBondageItem.createStack(
            new ResourceLocation("tiedup", "classic_earplugs")
        );
    }

    // COLOR METHODS

    /** NBT key for CustomModelData (used for model overrides) */
    public static final String NBT_CUSTOM_MODEL_DATA = "CustomModelData";

    /**
     * Apply color NBT to an ItemStack.
     * Sets both the ItemColor name and CustomModelData for model selection.
     */
    public static void applyColor(ItemStack stack, ItemColor color) {
        if (stack.isEmpty() || color == null) return;
        CompoundTag tag = stack.getOrCreateTag();
        tag.putString(NBT_ITEM_COLOR, color.getName());
        tag.putInt(NBT_CUSTOM_MODEL_DATA, color.getModelId());
    }

    /**
     * Get color from an ItemStack.
     * @return The color, or null if no color is set
     */
    @Nullable
    public static ItemColor getColor(ItemStack stack) {
        if (stack.isEmpty()) return null;
        CompoundTag tag = stack.getTag();
        if (tag == null || !tag.contains(NBT_ITEM_COLOR)) return null;
        return ItemColor.fromName(tag.getString(NBT_ITEM_COLOR));
    }

    /**
     * Check if an ItemStack has a color applied.
     */
    public static boolean hasColor(ItemStack stack) {
        return getColor(stack) != null;
    }

    /**
     * Get the texture suffix for an item's color.
     * Example: "ropes" + "_red" = "ropes_red"
     * @return The color suffix (e.g., "_red"), or empty string if no color
     */
    public static String getColorSuffix(ItemStack stack) {
        ItemColor color = getColor(stack);
        return color != null ? "_" + color.getName() : "";
    }

    // COLOR VALIDATION

    /**
     * Get a random color that has a texture for the given bind variant.
     * Excludes colors that don't have textures for specific variants.
     */
    public static ItemColor getValidColorForBind(BindVariant variant) {
        ItemColor color;
        int attempts = 0;
        do {
            color = ItemColor.getRandomStandard();
            attempts++;
            // Prevent infinite loop
            if (attempts > 50) break;
        } while (!isColorValidForBind(color, variant));
        return color;
    }

    /**
     * Check if a color has a texture for the given bind variant.
     * Returns false for colors without textures.
     */
    public static boolean isColorValidForBind(
        ItemColor color,
        BindVariant variant
    ) {
        if (color == null || variant == null) return true;

        // BROWN doesn't have textures for ROPES and SHIBARI
        if (
            color == ItemColor.BROWN &&
            (variant == BindVariant.ROPES || variant == BindVariant.SHIBARI)
        ) {
            return false;
        }

        // GRAY doesn't have texture for DUCT_TAPE
        if (color == ItemColor.GRAY && variant == BindVariant.DUCT_TAPE) {
            return false;
        }

        return true;
    }

    /**
     * Check if a color has a texture for the given gag variant.
     */
    public static boolean isColorValidForGag(
        ItemColor color,
        GagVariant variant
    ) {
        if (color == null || variant == null) return true;

        // GRAY doesn't have texture for TAPE_GAG
        if (color == ItemColor.GRAY && variant == GagVariant.TAPE_GAG) {
            return false;
        }

        // WHITE doesn't have texture for CLOTH_GAG and CLEAVE_GAG
        if (
            color == ItemColor.WHITE &&
            (variant == GagVariant.CLOTH_GAG ||
                variant == GagVariant.CLEAVE_GAG)
        ) {
            return false;
        }

        // RED doesn't have texture for BALL_GAG and BALL_GAG_STRAP
        if (
            color == ItemColor.RED &&
            (variant == GagVariant.BALL_GAG ||
                variant == GagVariant.BALL_GAG_STRAP)
        ) {
            return false;
        }

        return true;
    }

    /**
     * Check if a color has a texture for the given blindfold variant.
     */
    public static boolean isColorValidForBlindfold(
        ItemColor color,
        BlindfoldVariant variant
    ) {
        if (color == null || variant == null) return true;

        // BLACK doesn't have texture for CLASSIC or MASK blindfolds
        if (
            color == ItemColor.BLACK &&
            (variant == BlindfoldVariant.CLASSIC ||
                variant == BlindfoldVariant.MASK)
        ) {
            return false;
        }

        return true;
    }

    /**
     * Get a random color that has a texture for the given gag variant.
     */
    public static ItemColor getValidColorForGag(GagVariant variant) {
        ItemColor color;
        int attempts = 0;
        do {
            color = ItemColor.getRandomStandard();
            attempts++;
            if (attempts > 50) break;
        } while (!isColorValidForGag(color, variant));
        return color;
    }

    /**
     * Get a random color that has a texture for the given blindfold variant.
     */
    public static ItemColor getValidColorForBlindfold(
        BlindfoldVariant variant
    ) {
        ItemColor color;
        int attempts = 0;
        do {
            color = ItemColor.getRandomStandard();
            attempts++;
            if (attempts > 50) break;
        } while (!isColorValidForBlindfold(color, variant));
        return color;
    }
}