package com.tiedup.remake.v2.bondage.datadriven;

import com.google.gson.JsonObject;
import com.tiedup.remake.v2.bondage.component.ComponentHolder;
import com.tiedup.remake.v2.bondage.component.ComponentType;
import com.tiedup.remake.v2.bondage.component.IItemComponent;
import java.util.Collection;
import java.util.Collections;
import java.util.EnumMap;
import java.util.HashMap;
import java.util.Map;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.item.ItemStack;
import org.jetbrains.annotations.Nullable;

/**
 * Thread-safe registry for data-driven bondage item definitions.
 *
 * <p>Populated by the reload listener that scans {@code tiedup_items/} JSON files.
 * Uses a single volatile snapshot reference to ensure readers always see a
 * consistent pair of definitions + component holders. This prevents torn reads
 * where one map is updated but the other is stale.</p>
 *
 * <p>Lookup methods accept either a {@link ResourceLocation} ID directly
 * or an {@link ItemStack} (reads the {@code tiedup_item_id} NBT tag).</p>
 */
public final class DataDrivenItemRegistry {

    /** NBT key storing the data-driven item ID on ItemStacks. */
    public static final String NBT_ITEM_ID = "tiedup_item_id";

    /**
     * Immutable snapshot combining definitions and their component holders.
     * Swapped atomically via a single volatile write to prevent torn reads.
     */
    private record RegistrySnapshot(
        Map<ResourceLocation, DataDrivenItemDefinition> definitions,
        Map<ResourceLocation, ComponentHolder> holders
    ) {
        static final RegistrySnapshot EMPTY = new RegistrySnapshot(Map.of(), Map.of());
    }

    /**
     * Single volatile reference to the current registry state.
     * All read methods capture this reference ONCE at the start to ensure
     * consistency within a single call.
     */
    private static volatile RegistrySnapshot SNAPSHOT = RegistrySnapshot.EMPTY;

    /**
     * Monotonically increasing revision counter, incremented on every
     * {@link #reload} or {@link #mergeAll}. Client-side caches (e.g.,
     * {@code DataDrivenIconOverrides}) observe this to invalidate themselves
     * when the registry changes between bake passes — {@code /reload} updates
     * the registry but does not re-fire {@code ModelEvent.ModifyBakingResult},
     * so any cached resolution would otherwise go stale.
     */
    private static volatile int revision = 0;

    /** Guards read-then-write sequences in {@link #reload} and {@link #mergeAll}. */
    private static final Object RELOAD_LOCK = new Object();

    private DataDrivenItemRegistry() {}

    /**
     * Atomically replace all definitions with a new set.
     * Called by the reload listener after parsing all JSON files.
     *
     * @param newDefs the new definitions map (will be defensively copied)
     */
    public static void reload(
        Map<ResourceLocation, DataDrivenItemDefinition> newDefs
    ) {
        synchronized (RELOAD_LOCK) {
            Map<ResourceLocation, DataDrivenItemDefinition> defs =
                Collections.unmodifiableMap(new HashMap<>(newDefs));
            SNAPSHOT = new RegistrySnapshot(defs, buildComponentHolders(defs));
            revision++;
        }
    }

    /**
     * Atomically merge new definitions into the existing registry.
     *
     * <p>On an integrated server, both the client (assets/) and server (data/) reload
     * listeners populate this registry. Using {@link #reload} would cause the second
     * listener to overwrite the first's definitions. This method builds a new map
     * from the existing snapshot + the new entries, then swaps atomically as a
     * single snapshot to prevent torn reads.</p>
     *
     * @param newDefs the definitions to merge (will overwrite existing entries with same key)
     */
    public static void mergeAll(
        Map<ResourceLocation, DataDrivenItemDefinition> newDefs
    ) {
        synchronized (RELOAD_LOCK) {
            Map<ResourceLocation, DataDrivenItemDefinition> merged = new HashMap<>(
                SNAPSHOT.definitions
            );
            merged.putAll(newDefs);
            Map<ResourceLocation, DataDrivenItemDefinition> defs =
                Collections.unmodifiableMap(merged);
            SNAPSHOT = new RegistrySnapshot(defs, buildComponentHolders(defs));
            revision++;
        }
    }

    /**
     * Current registry revision — incremented on every {@link #reload} or
     * {@link #mergeAll}. Client-side caches compare their stored revision
     * against this to detect staleness across {@code /reload} cycles.
     */
    public static int getRevision() {
        return revision;
    }

    /**
     * Get a definition by its unique ID.
     *
     * @param id the definition ID (e.g., "tiedup:leather_armbinder")
     * @return the definition, or null if not found
     */
    @Nullable
    public static DataDrivenItemDefinition get(ResourceLocation id) {
        RegistrySnapshot snap = SNAPSHOT;
        return snap.definitions.get(id);
    }

    /**
     * Get a definition from an ItemStack by reading the {@code tiedup_item_id} NBT tag.
     *
     * @param stack the ItemStack to inspect
     * @return the definition, or null if the stack is empty, has no tag, or the ID is unknown
     */
    @Nullable
    public static DataDrivenItemDefinition get(ItemStack stack) {
        if (stack.isEmpty()) return null;
        CompoundTag tag = stack.getTag();
        if (tag == null || !tag.contains(NBT_ITEM_ID)) return null;
        ResourceLocation id = ResourceLocation.tryParse(
            tag.getString(NBT_ITEM_ID)
        );
        if (id == null) return null;
        RegistrySnapshot snap = SNAPSHOT;
        return snap.definitions.get(id);
    }

    /**
     * Get all registered definitions.
     *
     * @return unmodifiable collection of all definitions
     */
    public static Collection<DataDrivenItemDefinition> getAll() {
        RegistrySnapshot snap = SNAPSHOT;
        return snap.definitions.values();
    }

    /**
     * Clear all definitions. Called on world unload or for testing.
     */
    public static void clear() {
        synchronized (RELOAD_LOCK) {
            SNAPSHOT = RegistrySnapshot.EMPTY;
        }
    }

    // ===== COMPONENT HOLDERS =====

    /**
     * Get the ComponentHolder for a data-driven item stack.
     *
     * <p>Captures the snapshot reference once to ensure consistent reads
     * between the definition lookup and the holder lookup.</p>
     *
     * @param stack the ItemStack to inspect
     * @return the holder, or null if the stack is not data-driven or has no definition
     */
    @Nullable
    public static ComponentHolder getComponents(ItemStack stack) {
        if (stack.isEmpty()) return null;
        CompoundTag tag = stack.getTag();
        if (tag == null || !tag.contains(NBT_ITEM_ID)) return null;
        ResourceLocation id = ResourceLocation.tryParse(
            tag.getString(NBT_ITEM_ID)
        );
        if (id == null) return null;
        // Capture snapshot once to ensure definition and holder come from
        // the same atomic snapshot, preventing torn reads.
        RegistrySnapshot snap = SNAPSHOT;
        DataDrivenItemDefinition def = snap.definitions.get(id);
        if (def == null) return null;
        return snap.holders.get(def.id());
    }

    /**
     * Get the ComponentHolder for a data-driven item by its definition ID.
     *
     * @param id the definition ID
     * @return the holder, or null if not found
     */
    @Nullable
    public static ComponentHolder getComponents(ResourceLocation id) {
        RegistrySnapshot snap = SNAPSHOT;
        return snap.holders.get(id);
    }

    /**
     * Build component holders from a definitions map.
     * Each definition's raw componentConfigs are instantiated via
     * {@link ComponentType#create(JsonObject)}.
     */
    private static Map<ResourceLocation, ComponentHolder> buildComponentHolders(
        Map<ResourceLocation, DataDrivenItemDefinition> definitions
    ) {
        Map<ResourceLocation, ComponentHolder> holders = new HashMap<>();
        for (Map.Entry<ResourceLocation, DataDrivenItemDefinition> entry :
            definitions.entrySet()) {
            DataDrivenItemDefinition def = entry.getValue();
            Map<ComponentType, IItemComponent> components =
                new EnumMap<>(ComponentType.class);
            for (Map.Entry<ComponentType, JsonObject> compEntry :
                def.componentConfigs().entrySet()) {
                components.put(
                    compEntry.getKey(),
                    compEntry.getKey().create(compEntry.getValue())
                );
            }
            holders.put(
                entry.getKey(),
                components.isEmpty()
                    ? ComponentHolder.EMPTY
                    : new ComponentHolder(components)
            );
        }
        return Collections.unmodifiableMap(holders);
    }
}
