TiedUp-/src/main/java/com/tiedup/remake/client/gltf/GltfAnimationApplier.java

package com.tiedup.remake.client.gltf;

import com.tiedup.remake.client.animation.BondageAnimationManager;
import com.tiedup.remake.client.animation.context.AnimationContext;
import com.tiedup.remake.client.animation.context.ContextAnimationFactory;
import com.tiedup.remake.client.animation.context.GlbAnimationResolver;
import com.tiedup.remake.client.animation.context.RegionBoneMapper;
import dev.kosmx.playerAnim.core.data.KeyframeAnimation;
import java.util.Collections;
import java.util.HashSet;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.entity.LivingEntity;
import net.minecraftforge.api.distmarker.Dist;
import net.minecraftforge.api.distmarker.OnlyIn;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.jetbrains.annotations.Nullable;

/**
 * V2 Animation Applier -- manages dual-layer animation for V2 bondage items.
 *
 * <p>Orchestrates two PlayerAnimator layers simultaneously:
 * <ul>
 *   <li><b>Context layer</b> (priority 40): base body posture (stand/sit/kneel/sneak/walk)
 *       with item-owned parts disabled, via {@link ContextAnimationFactory}</li>
 *   <li><b>Item layer</b> (priority 42): per-item GLB animation with only owned bones enabled,
 *       via {@link GltfPoseConverter#convertSelective}</li>
 * </ul>
 *
 * <p>Each equipped V2 item controls ONLY the bones matching its occupied body regions.
 * Bones not owned by any item pass through from the context layer, which provides the
 * appropriate base posture animation.
 *
 * <p>State tracking avoids redundant animation replays: a composite key of
 * {@code animSource|context|ownedParts} is compared per-entity to skip no-op updates.
 *
 * <p>Item animations are cached by {@code animSource#context#ownedParts} since the same
 * GLB + context + owned parts always produces the same KeyframeAnimation.
 *
 * @see ContextAnimationFactory
 * @see GlbAnimationResolver
 * @see GltfPoseConverter#convertSelective
 * @see BondageAnimationManager
 */
@OnlyIn(Dist.CLIENT)
public final class GltfAnimationApplier {

    private static final Logger LOGGER = LogManager.getLogger("GltfPipeline");

    /**
     * Cache of converted item-layer KeyframeAnimations, keyed by
     * {@code animSource#context#ownedPartsHash}. LRU-bounded via
     * access-ordered {@link LinkedHashMap}: size capped at
     * {@link #ITEM_ANIM_CACHE_MAX}, head (least-recently-used) evicted on
     * overflow. Wrapped in {@link Collections#synchronizedMap} because
     * {@code LinkedHashMap.get} mutates the iteration order. External
     * iteration is not supported; use only {@code get}/{@code put}/{@code clear}.
     */
    private static final int ITEM_ANIM_CACHE_MAX = 256;

    // Initial capacity (int)(cap / loadFactor) + 1 so the cap is reached
    // without a rehash.
    private static final int ITEM_ANIM_CACHE_INITIAL_CAPACITY =
        (int) (ITEM_ANIM_CACHE_MAX / 0.75f) + 1;

    private static final Map<String, KeyframeAnimation> itemAnimCache =
        Collections.synchronizedMap(
            new LinkedHashMap<String, KeyframeAnimation>(
                ITEM_ANIM_CACHE_INITIAL_CAPACITY,
                0.75f,
                true // access-order
            ) {
                @Override
                protected boolean removeEldestEntry(
                    Map.Entry<String, KeyframeAnimation> eldest
                ) {
                    return size() > ITEM_ANIM_CACHE_MAX;
                }
            }
        );

    /**
     * Track which composite state is currently active per entity, to avoid redundant replays.
     * Keyed by entity UUID, value is "animSource|context|sortedParts".
     */
    private static final Map<UUID, String> activeStateKeys =
        new ConcurrentHashMap<>();

    /** Track cache keys where GLB loading failed, to avoid per-tick retries. */
    private static final Set<String> failedLoadKeys =
        ConcurrentHashMap.newKeySet();

    private GltfAnimationApplier() {}

    // INIT (legacy)

    /**
     * Legacy init method -- called by GltfClientSetup.
     * No-op: layer registration is handled by {@link BondageAnimationManager#init()}.
     */
    public static void init() {
        // No-op: animation layers are managed by BondageAnimationManager
    }

    // V2 DUAL-LAYER API

    /**
     * Apply the full V2 animation state: context layer + item layer.
     *
     * <p>Flow:
     * <ol>
     *   <li>Build a composite state key and skip if unchanged</li>
     *   <li>Create/retrieve a context animation with disabledOnContext parts disabled,
     *       play on context layer via {@link BondageAnimationManager#playContext}</li>
     *   <li>Load the GLB (from {@code animationSource} or {@code modelLoc}),
     *       resolve the named animation via {@link GlbAnimationResolver#resolve},
     *       convert with selective parts via {@link GltfPoseConverter#convertSelective},
     *       play on item layer via {@link BondageAnimationManager#playDirect}</li>
     * </ol>
     *
     * <p>The ownership model enables "free bone" animation: if a bone is not claimed
     * by any item, the winning item can animate it IF its GLB has keyframes for that bone.
     * This allows a straitjacket (ARMS+TORSO) to also animate free legs.</p>
     *
     * @param entity          the entity to animate
     * @param modelLoc        the item's GLB model (for mesh rendering, and default animation source)
     * @param animationSource separate GLB for animations (shared template), or null to use modelLoc
     * @param context         current animation context (STAND_IDLE, SIT_IDLE, etc.)
     * @param ownership       bone ownership: which parts this item owns vs other items
     * @return true if the item layer animation was applied successfully
     */
    public static boolean applyV2Animation(
        LivingEntity entity,
        ResourceLocation modelLoc,
        @Nullable ResourceLocation animationSource,
        AnimationContext context,
        RegionBoneMapper.BoneOwnership ownership
    ) {
        if (entity == null || modelLoc == null) return false;

        ResourceLocation animSource =
            animationSource != null ? animationSource : modelLoc;
        // Cache key includes both owned and enabled parts for full disambiguation
        String ownedKey = canonicalPartsKey(ownership.thisParts());
        String enabledKey = canonicalPartsKey(ownership.enabledParts());
        String partsKey = ownedKey + ";" + enabledKey;

        // Build composite state key to detect context changes.
        // NOTE: This key does NOT include the variant name — that is resolved fresh
        // each time the context changes, enabling random variant selection.
        String stateKey = animSource + "|" + context.name() + "|" + partsKey;
        String currentKey = activeStateKeys.get(entity.getUUID());
        if (stateKey.equals(currentKey)) {
            return true; // Same context, same parts — no need to re-resolve
        }

        // === Layer 1: Context animation (base body posture) ===
        // Parts owned by ANY item (this or others) are disabled on the context layer.
        // Only free parts remain enabled on context.
        KeyframeAnimation contextAnim = ContextAnimationFactory.create(
            context,
            ownership.disabledOnContext()
        );
        if (contextAnim != null) {
            BondageAnimationManager.playContext(entity, contextAnim);
        }

        // === Layer 2: Item animation (GLB pose with selective bones) ===
        String itemCacheKey = buildItemCacheKey(animSource, context, partsKey);

        // Skip if this GLB already failed to load
        if (failedLoadKeys.contains(itemCacheKey)) {
            activeStateKeys.put(entity.getUUID(), stateKey);
            return false;
        }

        // Resolve animation data first (needed for variant resolution)
        GltfData animData = GlbAnimationResolver.resolveAnimationData(
            modelLoc,
            animationSource
        );
        if (animData == null) {
            LOGGER.warn(
                "[GltfPipeline] Failed to load animation GLB: {}",
                animSource
            );
            failedLoadKeys.add(itemCacheKey);
            activeStateKeys.put(entity.getUUID(), stateKey);
            return false;
        }

        // Resolve which named animation to use (with fallback chain + variant selection).
        // This must happen BEFORE the cache lookup because variant selection is random —
        // we want a fresh random pick each time the context changes, not a permanently
        // cached first pick.
        String glbAnimName = GlbAnimationResolver.resolve(
            animData,
            context
        );

        // Include the resolved animation name in the cache key so different variants
        // (Struggle.1 vs Struggle.2) get separate cache entries.
        String variantCacheKey = itemCacheKey + "#" + (glbAnimName != null ? glbAnimName : "default");

        // Atomic get-or-compute under the map's monitor. Collections
        // .synchronizedMap only synchronizes individual get/put calls, so a
        // naive check-then-put races between concurrent converters and can
        // both double-convert and trip removeEldestEntry with a stale size.
        KeyframeAnimation itemAnim;
        synchronized (itemAnimCache) {
            itemAnim = itemAnimCache.get(variantCacheKey);
            if (itemAnim == null) {
                itemAnim = GltfPoseConverter.convertSelective(
                    animData,
                    glbAnimName,
                    ownership.thisParts(),
                    ownership.enabledParts()
                );
                itemAnimCache.put(variantCacheKey, itemAnim);
            }
        }

        BondageAnimationManager.playDirect(entity, itemAnim);

        activeStateKeys.put(entity.getUUID(), stateKey);
        return true;
    }

    /**
     * Apply V2 animation from ALL equipped items simultaneously.
     *
     * <p>Each item contributes keyframes for only its owned bones into a shared
     * {@link KeyframeAnimation.AnimationBuilder}. The first item in the list (highest priority)
     * can additionally animate free bones if its GLB has keyframes for them.</p>
     *
     * @param entity       the entity to animate
     * @param items        resolved V2 items with per-item ownership, sorted by priority desc
     * @param context      current animation context
     * @param allOwnedParts union of all owned parts across all items
     * @return true if the composite animation was applied
     */
    public static boolean applyMultiItemV2Animation(
        LivingEntity entity,
        List<RegionBoneMapper.V2ItemAnimInfo> items,
        AnimationContext context,
        Set<String> allOwnedParts
    ) {
        if (entity == null || items.isEmpty()) return false;

        // Build composite state key
        StringBuilder keyBuilder = new StringBuilder();
        for (RegionBoneMapper.V2ItemAnimInfo item : items) {
            ResourceLocation src =
                item.animSource() != null ? item.animSource() : item.modelLoc();
            keyBuilder
                .append(src)
                .append(':')
                .append(canonicalPartsKey(item.ownedParts()))
                .append(';');
        }
        keyBuilder.append(context.name());
        String stateKey = keyBuilder.toString();

        String currentKey = activeStateKeys.get(entity.getUUID());
        if (stateKey.equals(currentKey)) {
            return true; // Already active
        }

        // === Layer 1: Context animation ===
        KeyframeAnimation contextAnim = ContextAnimationFactory.create(
            context,
            allOwnedParts
        );
        if (contextAnim != null) {
            BondageAnimationManager.playContext(entity, contextAnim);
        }

        // === Layer 2: Composite item animation ===
        // Pre-resolve animation data and variant names for all items BEFORE cache lookup.
        // This ensures random variant selection happens fresh on each context change,
        // and each variant combination gets its own cache entry.
        record ResolvedItem(
            GltfData animData,
            String glbAnimName,
            RegionBoneMapper.V2ItemAnimInfo info
        ) {}

        List<ResolvedItem> resolvedItems = new ArrayList<>();
        StringBuilder variantKeyBuilder = new StringBuilder("multi#").append(stateKey);

        for (RegionBoneMapper.V2ItemAnimInfo item : items) {
            ResourceLocation animSource =
                item.animSource() != null ? item.animSource() : item.modelLoc();

            GltfData animData = GlbAnimationResolver.resolveAnimationData(
                item.modelLoc(), item.animSource()
            );
            if (animData == null) {
                LOGGER.warn(
                    "[GltfPipeline] Failed to load GLB for multi-item: {}",
                    animSource
                );
                continue;
            }

            String glbAnimName = GlbAnimationResolver.resolve(animData, context);
            resolvedItems.add(new ResolvedItem(animData, glbAnimName, item));
            variantKeyBuilder.append('#')
                .append(glbAnimName != null ? glbAnimName : "default");
        }

        String compositeCacheKey = variantKeyBuilder.toString();

        if (failedLoadKeys.contains(compositeCacheKey)) {
            activeStateKeys.put(entity.getUUID(), stateKey);
            return false;
        }

        // Atomic get-or-compute under the map's monitor (see
        // applyV2Animation). All current callers are render-thread so no
        // contention in practice, but the synchronized wrap closes the
        // window where two converters could race and clobber each other.
        KeyframeAnimation compositeAnim;
        synchronized (itemAnimCache) {
            compositeAnim = itemAnimCache.get(compositeCacheKey);
        }
        if (compositeAnim == null) {
            KeyframeAnimation.AnimationBuilder builder =
                new KeyframeAnimation.AnimationBuilder(
                    dev.kosmx.playerAnim.core.data.AnimationFormat.JSON_EMOTECRAFT
                );
            builder.beginTick = 0;
            builder.isLooped = true;
            builder.returnTick = 0;
            builder.name = "gltf_composite";

            boolean anyLoaded = false;
            int maxEndTick = 1;
            Set<String> unionKeyframeParts = new HashSet<>();
            boolean anyFullBody = false;
            boolean anyFullHead = false;

            for (ResolvedItem resolved : resolvedItems) {
                RegionBoneMapper.V2ItemAnimInfo item = resolved.info();
                GltfData animData = resolved.animData();
                String glbAnimName = resolved.glbAnimName();
                ResourceLocation animSource =
                    item.animSource() != null ? item.animSource() : item.modelLoc();

                GltfData.AnimationClip rawClip;
                if (glbAnimName != null) {
                    rawClip = animData.getRawAnimation(glbAnimName);
                } else {
                    rawClip = null;
                }
                if (rawClip == null) {
                    rawClip = animData.rawGltfAnimation();
                }

                // Compute effective parts: intersect animation_bones whitelist with ownedParts
                // if the item declares per-animation bone filtering.
                Set<String> effectiveParts = item.ownedParts();
                if (glbAnimName != null && !item.animationBones().isEmpty()) {
                    Set<String> override = item.animationBones().get(glbAnimName);
                    if (override != null) {
                        Set<String> filtered = new HashSet<>(override);
                        filtered.retainAll(item.ownedParts());
                        if (!filtered.isEmpty()) {
                            effectiveParts = filtered;
                        }
                    }
                }

                Set<String> itemKeyframeParts =
                    GltfPoseConverter.addBonesToBuilder(
                        builder, animData, rawClip, effectiveParts
                    );
                unionKeyframeParts.addAll(itemKeyframeParts);
                maxEndTick = Math.max(
                    maxEndTick,
                    GltfPoseConverter.computeEndTick(rawClip)
                );
                // FullX / FullHeadX opt-in: ANY item requesting it lifts the
                // restriction for the composite. The animation name passed to
                // the core helper uses the same "gltf_" prefix convention as
                // the single-item path.
                String prefixed = glbAnimName != null
                    ? "gltf_" + glbAnimName
                    : null;
                if (GltfPoseConverter.isFullBodyAnimName(prefixed)) {
                    anyFullBody = true;
                }
                if (GltfPoseConverter.isFullHeadAnimName(prefixed)) {
                    anyFullHead = true;
                }
                anyLoaded = true;

                LOGGER.debug(
                    "[GltfPipeline] Multi-item: {} -> owned={}, effective={}, anim={}",
                    animSource, item.ownedParts(), effectiveParts, glbAnimName
                );
            }

            if (!anyLoaded) {
                failedLoadKeys.add(compositeCacheKey);
                activeStateKeys.put(entity.getUUID(), stateKey);
                return false;
            }

            builder.endTick = maxEndTick;
            builder.stopTick = maxEndTick;

            // Selective-part enabling for the composite. Owned parts always on;
            // free parts (including head) opt-in only if ANY item declares a
            // FullX / FullHeadX animation AND has keyframes for that part.
            GltfPoseConverter.enableSelectivePartsComposite(
                builder,
                allOwnedParts,
                unionKeyframeParts,
                anyFullBody,
                anyFullHead
            );

            compositeAnim = builder.build();
            synchronized (itemAnimCache) {
                // Another thread may have computed the same key while we were
                // building. Prefer its result to keep one instance per key,
                // matching removeEldestEntry's accounting.
                KeyframeAnimation winner = itemAnimCache.get(compositeCacheKey);
                if (winner != null) {
                    compositeAnim = winner;
                } else {
                    itemAnimCache.put(compositeCacheKey, compositeAnim);
                }
            }
        }

        BondageAnimationManager.playDirect(entity, compositeAnim);
        activeStateKeys.put(entity.getUUID(), stateKey);
        return true;
    }

    // CLEAR / QUERY

    /**
     * Clear all V2 animation layers from an entity and remove tracking.
     * Stops both the context layer and the item layer.
     *
     * @param entity the entity to clear animations from
     */
    public static void clearV2Animation(LivingEntity entity) {
        if (entity == null) return;
        activeStateKeys.remove(entity.getUUID());
        BondageAnimationManager.stopContext(entity);
        BondageAnimationManager.stopAnimation(entity);
    }

    /**
     * Check if an entity has active V2 animation state.
     *
     * @param entity the entity to check
     * @return true if the entity has an active V2 animation state key
     */
    public static boolean hasActiveState(LivingEntity entity) {
        return entity != null && activeStateKeys.containsKey(entity.getUUID());
    }

    /**
     * Remove tracking for an entity (e.g., on logout/unload).
     * Does NOT stop any currently playing animation -- use {@link #clearV2Animation} for that.
     *
     * @param entityId UUID of the entity to stop tracking
     */
    public static void removeTracking(UUID entityId) {
        activeStateKeys.remove(entityId);
    }

    // CACHE MANAGEMENT

    /**
     * Invalidate all cached item animations and tracking state.
     * Call this on resource reload (F3+T) to pick up changed GLB/JSON files.
     *
     * <p>Does NOT clear ContextAnimationFactory or ContextGlbRegistry here.
     * Those are cleared in the reload listener AFTER ContextGlbRegistry.reload()
     * to prevent the render thread from caching stale JSON fallbacks during
     * the window between clear and repopulate.</p>
     */
    public static void invalidateCache() {
        itemAnimCache.clear();
        activeStateKeys.clear();
        failedLoadKeys.clear();
    }

    /**
     * Clear all state (cache + tracking). Called on world unload.
     * Clears everything including context caches (no concurrent reload during unload).
     */
    public static void clearAll() {
        itemAnimCache.clear();
        activeStateKeys.clear();
        failedLoadKeys.clear();
        com.tiedup.remake.client.animation.context.ContextGlbRegistry.clear();
        ContextAnimationFactory.clearCache();
    }

    // INTERNAL

    /**
     * Build cache key for item-layer animations.
     * Format: "animSource#contextName#sortedParts"
     */
    private static String buildItemCacheKey(
        ResourceLocation animSource,
        AnimationContext context,
        String partsKey
    ) {
        return animSource + "#" + context.name() + "#" + partsKey;
    }

    /**
     * Build a canonical, deterministic string from the owned parts set.
     * Sorted alphabetically and joined by comma — guarantees no hash collisions.
     */
    private static String canonicalPartsKey(Set<String> ownedParts) {
        return String.join(",", new TreeSet<>(ownedParts));
    }

}