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.
 *
 * All bondage items are now created via DataDrivenBondageItem.createStack().
 *
 * 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;
        }

        public boolean hasGag() { return !gag.isEmpty(); }
        public boolean hasMittens() { return !mittens.isEmpty(); }
        public boolean hasEarplugs() { return !earplugs.isEmpty(); }
        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);
    }

    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;
    }

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

        // 2. Select color (if theme supports it)
        ItemColor color = theme.supportsColor()
            ? ItemColor.getRandomStandard()
            : null;

        // 3. Create bind (always)
        ItemStack bind = createItemById(theme.getBindId(), color);

        // 4. Roll for gag
        ItemStack gag = ItemStack.EMPTY;
        double gagProb = getAdjustedProbability(
            PROB_GAG, ELITE_GAG_BONUS, ARCHER_GAG_PENALTY, isElite, isArcher
        );
        if (RANDOM.nextDouble() < gagProb) {
            gag = createItemById(theme.getRandomGagId(), color);
        }

        // 5. Roll for mittens
        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
        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
        ItemStack blindfold = ItemStack.EMPTY;
        double blindfoldProb = getAdjustedProbability(
            PROB_BLINDFOLD, ELITE_BLINDFOLD_BONUS, ARCHER_BLINDFOLD_PENALTY, isElite, isArcher
        );
        if (theme.hasBlindfolds() && RANDOM.nextDouble() < blindfoldProb) {
            blindfold = createItemById(theme.getRandomBlindfoldId(), color);
        }

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

    // ITEM CREATION METHODS

    /**
     * Create a data-driven bondage item by registry name, with optional color.
     */
    public static ItemStack createItemById(String id, @Nullable ItemColor color) {
        ItemStack stack = DataDrivenBondageItem.createStack(
            new ResourceLocation("tiedup", id)
        );
        if (color != null) {
            applyColor(stack, color);
        }
        return stack;
    }

    /**
     * Create mittens ItemStack.
     */
    public static ItemStack createMittens() {
        return DataDrivenBondageItem.createStack(
            new ResourceLocation("tiedup", "leather_mittens")
        );
    }

    /**
     * Create earplugs ItemStack.
     */
    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.
     */
    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.
     */
    public static String getColorSuffix(ItemStack stack) {
        ItemColor color = getColor(stack);
        return color != null ? "_" + color.getName() : "";
    }
}