package com.tiedup.remake.v2.client;

import com.tiedup.remake.v2.bondage.datadriven.DataDrivenItemDefinition;
import com.tiedup.remake.v2.bondage.datadriven.DataDrivenItemRegistry;
import com.tiedup.remake.v2.furniture.FurnitureDefinition;
import com.tiedup.remake.v2.furniture.FurniturePlacerItem;
import com.tiedup.remake.v2.furniture.FurnitureRegistry;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import net.minecraft.client.Minecraft;
import net.minecraft.client.multiplayer.ClientLevel;
import net.minecraft.client.renderer.block.model.ItemOverrides;
import net.minecraft.client.resources.model.BakedModel;
import net.minecraft.client.resources.model.ModelResourceLocation;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.item.ItemStack;
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;

/**
 * Custom {@link ItemOverrides} that switches the rendered item model based on
 * NBT-driven icon definitions.
 *
 * <p>Both {@link com.tiedup.remake.v2.bondage.datadriven.DataDrivenBondageItem} and
 * {@link com.tiedup.remake.v2.furniture.FurniturePlacerItem} are singleton Items
 * where each stack's behavior varies by NBT. This override reads the item/furniture
 * ID from NBT, looks up the corresponding definition's {@code icon} field, and
 * resolves the matching {@link BakedModel} from the model registry.</p>
 *
 * <p>Icon model resolution strategy (tried in order):
 * <ol>
 *   <li>Plain {@link ResourceLocation} lookup (for models registered via
 *       {@link net.minecraftforge.client.event.ModelEvent.RegisterAdditional})</li>
 *   <li>{@link ModelResourceLocation} with "inventory" variant (for models baked
 *       from registered items, e.g., "tiedup:armbinder#inventory")</li>
 * </ol>
 * Falls back to the original model if no icon is defined or the icon model is not found.</p>
 */
@OnlyIn(Dist.CLIENT)
public class DataDrivenIconOverrides extends ItemOverrides {

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

    /**
     * Identifies which type of NBT-driven item this override handles.
     */
    public enum Mode {
        /** Data-driven bondage items (reads {@code tiedup_item_id} NBT). */
        BONDAGE_ITEM,
        /** Furniture placer items (reads {@code tiedup_furniture_id} NBT). */
        FURNITURE_PLACER,
    }

    private final Mode mode;

    /**
     * Sentinel meaning "no revision observed yet". Anything other than this
     * value means the caches were populated under that revision.
     */
    private static final int UNINITIALIZED = Integer.MIN_VALUE;

    /**
     * Last-observed revision of the backing registry (selected by {@link #mode}).
     * If the registry mutates between bake passes (typically via {@code /reload}
     * on the server, which does NOT re-fire {@code ModelEvent.ModifyBakingResult}),
     * the next {@link #resolve} call detects the mismatch and flushes
     * {@link #iconModelCache} / {@link #knownMissing} / {@link #warnedMissing}.
     *
     * <p>Lazily initialized on first {@link #resolve} to avoid a race where the
     * client reload listener runs after {@code ModifyBakingResult} created
     * this instance but before the first frame renders — early init here
     * would then observe a mismatch and flush an already-warm cache.</p>
     */
    private volatile int observedRevision = UNINITIALIZED;

    /**
     * Cache of resolved icon ResourceLocations to their BakedModels.
     * Cleared on resource reload (when ModifyBakingResult fires again) and
     * implicitly on registry revision change.
     * Uses ConcurrentHashMap because resolve() is called from the render thread.
     *
     * <p>Values are never null (ConcurrentHashMap forbids nulls). Missing icons
     * are tracked separately in {@link #knownMissing}.</p>
     */
    private final Map<ResourceLocation, BakedModel> iconModelCache =
        new ConcurrentHashMap<>();

    /**
     * Set of icon ResourceLocations that were looked up but not found.
     * Prevents repeated failed lookups from hitting the model registry every frame.
     */
    private final Map<ResourceLocation, Boolean> knownMissing =
        new ConcurrentHashMap<>();

    /**
     * Set of icon ResourceLocations that we already logged a warning for (to avoid log spam).
     */
    private final Map<ResourceLocation, Boolean> warnedMissing =
        new ConcurrentHashMap<>();

    public DataDrivenIconOverrides(Mode mode) {
        super();
        this.mode = mode;
    }

    @Override
    @Nullable
    public BakedModel resolve(
        BakedModel originalModel,
        ItemStack stack,
        @Nullable ClientLevel level,
        @Nullable LivingEntity entity,
        int seed
    ) {
        // Observe the registry revision; if it moved since we last populated
        // our caches, flush them. This handles the /reload-without-rebake path
        // where ModelEvent.ModifyBakingResult doesn't re-fire but the
        // underlying definition set has changed.
        //
        // The revision source depends on mode: bondage items live in
        // DataDrivenItemRegistry, furniture placers live in FurnitureRegistry.
        // A stale check against the wrong registry would ignore /reload updates
        // to the corresponding feature.
        int currentRevision = currentRegistryRevision();
        if (observedRevision == UNINITIALIZED) {
            // Lazy init on first resolve. Baseline the current revision so the
            // post-bake warm cache populated by the very first resolve call
            // isn't spuriously flushed by a race with the reload listener.
            observedRevision = currentRevision;
        } else if (currentRevision != observedRevision) {
            clearCache();
            observedRevision = currentRevision;
        }

        ResourceLocation iconRL = getIconFromStack(stack);
        if (iconRL == null) {
            // No icon defined for this variant — use the default model
            return originalModel;
        }

        // Check if already known to be missing
        if (knownMissing.containsKey(iconRL)) {
            return originalModel;
        }

        // Check cache for resolved model
        BakedModel cached = iconModelCache.get(iconRL);
        if (cached != null) {
            return cached;
        }

        // Try to resolve the icon model
        BakedModel resolved = lookupIconModel(iconRL);
        if (resolved != null) {
            iconModelCache.put(iconRL, resolved);
            return resolved;
        }

        // Mark as known missing to avoid repeated lookups
        knownMissing.put(iconRL, Boolean.TRUE);
        return originalModel;
    }

    /**
     * Read the icon ResourceLocation from the stack's NBT by looking up the definition.
     */
    @Nullable
    private ResourceLocation getIconFromStack(ItemStack stack) {
        if (stack.isEmpty()) return null;

        switch (mode) {
            case BONDAGE_ITEM: {
                DataDrivenItemDefinition def = DataDrivenItemRegistry.get(
                    stack
                );
                return def != null ? def.icon() : null;
            }
            case FURNITURE_PLACER: {
                String furnitureIdStr =
                    FurniturePlacerItem.getFurnitureIdFromStack(stack);
                if (furnitureIdStr == null) return null;
                FurnitureDefinition def = FurnitureRegistry.get(furnitureIdStr);
                return def != null ? def.icon() : null;
            }
            default:
                return null;
        }
    }

    /**
     * Look up an icon model from the model registry using multiple resolution strategies.
     *
     * <p>Strategy:
     * <ol>
     *   <li>Try the icon as a plain ResourceLocation (for RegisterAdditional models)</li>
     *   <li>If the icon path starts with "item/", strip it and try as a
     *       ModelResourceLocation with "inventory" variant</li>
     * </ol>
     *
     * @param iconRL the icon model ResourceLocation
     * @return the resolved BakedModel, or null if not found
     */
    @Nullable
    private BakedModel lookupIconModel(ResourceLocation iconRL) {
        Minecraft mc = Minecraft.getInstance();
        if (mc.getModelManager() == null) return null;

        BakedModel missingModel = mc.getModelManager().getMissingModel();

        // Strategy 1: Plain ResourceLocation lookup
        BakedModel model = mc.getModelManager().getModel(iconRL);
        if (model != missingModel) {
            return model;
        }

        // Strategy 2: Derive ModelResourceLocation with "inventory" variant
        // e.g., "tiedup:item/armbinder" → ModelResourceLocation("tiedup:armbinder", "inventory")
        String path = iconRL.getPath();
        if (path.startsWith("item/")) {
            String itemPath = path.substring("item/".length());
            ModelResourceLocation mrl = new ModelResourceLocation(
                ResourceLocation.fromNamespaceAndPath(iconRL.getNamespace(), itemPath),
                "inventory"
            );
            model = mc.getModelManager().getModel(mrl);
            if (model != missingModel) {
                return model;
            }
        }

        // Strategy 3: Try as-is with "inventory" variant
        // e.g., "tiedup:armbinder" → ModelResourceLocation("tiedup:armbinder", "inventory")
        if (!path.contains("/")) {
            ModelResourceLocation mrl = new ModelResourceLocation(
                iconRL,
                "inventory"
            );
            model = mc.getModelManager().getModel(mrl);
            if (model != missingModel) {
                return model;
            }
        }

        // Not found — log once
        if (!warnedMissing.containsKey(iconRL)) {
            warnedMissing.put(iconRL, Boolean.TRUE);
            LOGGER.warn(
                "[DataDrivenIcons] Icon model not found for '{}' (mode={}). " +
                    "Ensure a model JSON exists at assets/{}/models/{}.json or the item is registered.",
                iconRL,
                mode,
                iconRL.getNamespace(),
                iconRL.getPath()
            );
        }

        return null;
    }

    /**
     * Clear the icon model cache. Called when models are re-baked (on resource reload).
     */
    public void clearCache() {
        iconModelCache.clear();
        knownMissing.clear();
        warnedMissing.clear();
    }

    /**
     * Select the registry revision source appropriate for this override's mode.
     * Bondage items and furniture placers have independent registries and
     * independent /reload cycles, so a bondage-item cache must not invalidate
     * just because furniture reloaded (and vice versa).
     */
    private int currentRegistryRevision() {
        switch (mode) {
            case BONDAGE_ITEM:
                return DataDrivenItemRegistry.getRevision();
            case FURNITURE_PLACER:
                return com.tiedup.remake.v2.furniture.FurnitureRegistry.getRevision();
            default:
                return 0;
        }
    }
}
