package com.tiedup.remake.client.animation.context;

import com.mojang.logging.LogUtils;
import dev.kosmx.playerAnim.core.data.KeyframeAnimation;
import dev.kosmx.playerAnim.minecraftApi.PlayerAnimationRegistry;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import net.minecraft.resources.ResourceLocation;
import net.minecraftforge.api.distmarker.Dist;
import net.minecraftforge.api.distmarker.OnlyIn;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;

/**
 * Builds context {@link KeyframeAnimation}s with item-owned body parts disabled.
 *
 * <p>Context animations (loaded from {@code context_*.json} files in the PlayerAnimator
 * registry) control the base body posture -- standing, sitting, walking, etc.
 * When a V2 bondage item "owns" certain body parts (e.g., handcuffs own rightArm + leftArm),
 * those parts must NOT be driven by the context animation because the item's own
 * GLB animation controls them instead.</p>
 *
 * <p>This factory loads the base context animation, creates a mutable copy, disables
 * the owned parts, and builds an immutable result. Results are cached by
 * {@code contextSuffix|ownedPartsHash} to avoid repeated copies.</p>
 *
 * <p>Thread safety: the cache uses {@link ConcurrentHashMap}. All methods are
 * called from the render thread, but the concurrent map avoids issues if
 * resource reload triggers on a different thread.</p>
 *
 * @see AnimationContext
 * @see RegionBoneMapper#computeOwnedParts
 */
@OnlyIn(Dist.CLIENT)
public final class ContextAnimationFactory {

    private static final Logger LOGGER = LogUtils.getLogger();
    private static final String NAMESPACE = "tiedup";

    /**
     * Cache keyed by "contextSuffix|ownedPartsHashCode".
     * Null values are stored as sentinels for missing animations to avoid repeated lookups.
     */
    private static final Map<String, KeyframeAnimation> CACHE =
        new ConcurrentHashMap<>();

    /**
     * Sentinel set used to track cache keys where the base animation was not found,
     * so we don't log the same warning repeatedly.
     */
    private static final Set<String> MISSING_WARNED =
        ConcurrentHashMap.newKeySet();

    private ContextAnimationFactory() {}

    /**
     * Create (or retrieve from cache) a context animation with the given parts disabled.
     *
     * <p>If no parts need disabling, the base animation is returned as-is (no copy needed).
     * If the base animation is not found in the PlayerAnimator registry, returns null.</p>
     *
     * @param context       the current animation context (determines which context_*.json to load)
     * @param disabledParts set of PlayerAnimator part names to disable on the context layer
     *                      (e.g., {"rightArm", "leftArm"}), typically from
     *                      {@link RegionBoneMapper.BoneOwnership#disabledOnContext()}
     * @return the context animation with disabled parts suppressed, or null if not found
     */
    @Nullable
    public static KeyframeAnimation create(
        AnimationContext context,
        Set<String> disabledParts
    ) {
        String cacheKey =
            context.getAnimationSuffix() +
            "|" +
            String.join(",", new java.util.TreeSet<>(disabledParts));
        // computeIfAbsent cannot store null values, so we handle the missing case
        // by checking the MISSING_WARNED set to avoid redundant work.
        KeyframeAnimation cached = CACHE.get(cacheKey);
        if (cached != null) {
            return cached;
        }
        if (MISSING_WARNED.contains(cacheKey)) {
            return null;
        }

        KeyframeAnimation result = buildContextAnimation(
            context,
            disabledParts
        );
        if (result != null) {
            CACHE.put(cacheKey, result);
        } else {
            MISSING_WARNED.add(cacheKey);
        }
        return result;
    }

    /**
     * Build a context animation with the specified parts disabled.
     *
     * <p>Flow:
     * <ol>
     *   <li>Check {@link ContextGlbRegistry} for a GLB-based context animation (takes priority)</li>
     *   <li>Fall back to {@code tiedup:context_<suffix>} in PlayerAnimationRegistry (JSON-based)</li>
     *   <li>If no parts need disabling, return the base animation directly (immutable, shared)</li>
     *   <li>Otherwise, create a mutable copy via {@link KeyframeAnimation#mutableCopy()}</li>
     *   <li>Disable each part via {@link KeyframeAnimation.StateCollection#setEnabled(boolean)}</li>
     *   <li>Build and return the new immutable animation</li>
     * </ol>
     */
    @Nullable
    private static KeyframeAnimation buildContextAnimation(
        AnimationContext context,
        Set<String> disabledParts
    ) {
        String suffix = context.getAnimationSuffix();

        // Priority 1: GLB-based context animation from ContextGlbRegistry
        KeyframeAnimation baseAnim = ContextGlbRegistry.get(suffix);

        // Priority 2: JSON-based context animation from PlayerAnimationRegistry
        if (baseAnim == null) {
            ResourceLocation animId = ResourceLocation.fromNamespaceAndPath(
                NAMESPACE,
                "context_" + suffix
            );
            baseAnim = PlayerAnimationRegistry.getAnimation(animId);
        }

        if (baseAnim == null) {
            LOGGER.warn(
                "[V2Animation] Context animation not found for suffix: {}",
                suffix
            );
            return null;
        }

        if (disabledParts.isEmpty()) {
            return baseAnim;
        }

        // Create mutable copy so we can disable parts without affecting the registry/cache original
        KeyframeAnimation.AnimationBuilder builder = baseAnim.mutableCopy();
        disableParts(builder, disabledParts);
        return builder.build();
    }

    /**
     * Disable all animation axes on the specified parts.
     *
     * <p>Uses {@link KeyframeAnimation.AnimationBuilder#getPart(String)} to look up parts
     * by name, then {@link KeyframeAnimation.StateCollection#setEnabled(boolean)} to disable
     * all axes (x, y, z, pitch, yaw, roll, and bend/bendDirection if applicable).</p>
     *
     * <p>Unknown part names are silently ignored -- this can happen if the disabled parts set
     * includes future bone names not present in the current context animation.</p>
     */
    private static void disableParts(
        KeyframeAnimation.AnimationBuilder builder,
        Set<String> disabledParts
    ) {
        for (String partName : disabledParts) {
            KeyframeAnimation.StateCollection part = builder.getPart(partName);
            if (part != null) {
                part.setEnabled(false);
            }
        }
    }

    /**
     * Clear all cached animations. Call this on resource reload or when equipped items change
     * in a way that might invalidate cached part ownership.
     */
    public static void clearCache() {
        CACHE.clear();
        MISSING_WARNED.clear();
    }
}
