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

import com.tiedup.remake.v2.BodyRegionV2;
import com.tiedup.remake.v2.bondage.IV2BondageEquipment;
import com.tiedup.remake.v2.bondage.IV2BondageItem;
import java.util.Collections;
import java.util.EnumMap;
import java.util.IdentityHashMap;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import net.minecraft.core.BlockPos;
import net.minecraft.core.registries.Registries;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.ListTag;
import net.minecraft.nbt.StringTag;
import net.minecraft.nbt.Tag;
import net.minecraft.resources.ResourceKey;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.level.Level;
import org.jetbrains.annotations.Nullable;

/**
 * Concrete implementation of {@link IV2BondageEquipment}.
 *
 * Storage: EnumMap with 14 regions, each initialized to {@link ItemStack#EMPTY}.
 * Multi-region items share the same ItemStack reference across all occupied regions.
 *
 * NBT format:
 * - Root key: "V2BondageRegions" CompoundTag
 * - Per item: "REGION_NAME" -> ItemStack NBT
 * - Multi-region secondary slots: "REGION_NAME_also" -> ListTag of StringTag (other region names)
 * - Only non-empty regions are persisted
 */
public class V2BondageEquipment implements IV2BondageEquipment {

    private static final String NBT_ROOT_KEY = "V2BondageRegions";
    private static final String NBT_ALSO_SUFFIX = "_also";

    // Not thread-safe by itself. Safe in practice because Forge capabilities
    // maintain separate instances for logical client and logical server; each
    // side's instance is only touched by that side's main thread. If a future
    // feature ever reads/writes this from a worker (async asset, network
    // thread), the access pattern needs rethinking.
    private final EnumMap<BodyRegionV2, ItemStack> regions;

    // Pole leash persistence
    @Nullable
    private BlockPos savedPolePosition;

    @Nullable
    private ResourceKey<Level> savedPoleDimension;

    // Captor persistence
    @Nullable
    private UUID savedCaptorUUID;

    public V2BondageEquipment() {
        this.regions = new EnumMap<>(BodyRegionV2.class);
        for (BodyRegionV2 region : BodyRegionV2.values()) {
            regions.put(region, ItemStack.EMPTY);
        }
    }

    @Override
    public ItemStack getInRegion(BodyRegionV2 region) {
        if (region == null) return ItemStack.EMPTY;
        ItemStack stack = regions.get(region);
        return stack != null ? stack : ItemStack.EMPTY;
    }

    @Override
    public void setInRegion(BodyRegionV2 region, ItemStack stack) {
        if (region == null) return;
        regions.put(region, stack == null ? ItemStack.EMPTY : stack);
    }

    @Override
    public Map<BodyRegionV2, ItemStack> getAllEquipped() {
        // De-duplicate: multi-region items should appear only once.
        // Uses IdentityHashMap to track already-seen stack references.
        IdentityHashMap<ItemStack, BodyRegionV2> seen = new IdentityHashMap<>();
        Map<BodyRegionV2, ItemStack> result = new LinkedHashMap<>();

        for (BodyRegionV2 region : BodyRegionV2.values()) {
            ItemStack stack = regions.get(region);
            if (stack != null && !stack.isEmpty() && !seen.containsKey(stack)) {
                seen.put(stack, region);
                result.put(region, stack);
            }
        }

        return Collections.unmodifiableMap(result);
    }

    @Override
    public boolean isRegionOccupied(BodyRegionV2 region) {
        if (region == null) return false;
        ItemStack stack = regions.get(region);
        return stack != null && !stack.isEmpty();
    }

    /**
     * Check if a region is blocked by any equipped item's {@link IV2BondageItem#getBlockedRegions()}.
     * <p>
     * Scans all equipped items and returns true if any item blocks this region
     * (excluding self-blocking via the item's own occupied regions).
     * This replaces the old O(1) parent-global hierarchy check.
     */
    @Override
    public boolean isRegionBlocked(BodyRegionV2 region) {
        if (region == null) return false;
        // Scan regions.values() directly and identity-dedup multi-region
        // items inline. Avoids the 2-map allocation (IdentityHashMap +
        // LinkedHashMap) of getAllEquipped() — this is called per-frame
        // from FirstPersonHandHideHandler and would otherwise produce
        // 4 map allocs/frame even when the local player wears nothing.
        IdentityHashMap<ItemStack, Boolean> seen = null;
        for (ItemStack stack : regions.values()) {
            if (stack == null || stack.isEmpty()) continue;
            if (!(stack.getItem() instanceof IV2BondageItem item)) continue;
            if (seen == null) seen = new IdentityHashMap<>();
            if (seen.put(stack, Boolean.TRUE) != null) continue;
            if (
                item.getBlockedRegions(stack).contains(region) &&
                !item.getOccupiedRegions(stack).contains(region)
            ) {
                return true;
            }
        }
        return false;
    }

    @Override
    public int getEquippedCount() {
        // Count unique non-empty stacks directly, avoiding the 2-map allocation
        // of getAllEquipped(). Uses identity-based dedup for multi-region items.
        // Fast-path: zero allocations when no item is equipped (hot per-tick
        // call from hasAnyEquipment).
        IdentityHashMap<ItemStack, Boolean> seen = null;
        for (ItemStack stack : regions.values()) {
            if (stack != null && !stack.isEmpty()) {
                if (seen == null) seen = new IdentityHashMap<>();
                seen.put(stack, Boolean.TRUE);
            }
        }
        return seen == null ? 0 : seen.size();
    }

    @Override
    public void clearAll() {
        for (BodyRegionV2 region : BodyRegionV2.values()) {
            regions.put(region, ItemStack.EMPTY);
        }
    }

    // Pole leash persistence

    @Override
    public boolean wasLeashedToPole() {
        return savedPolePosition != null && savedPoleDimension != null;
    }

    @Override
    @Nullable
    public BlockPos getSavedPolePosition() {
        return savedPolePosition;
    }

    @Override
    @Nullable
    public ResourceKey<Level> getSavedPoleDimension() {
        return savedPoleDimension;
    }

    @Override
    public void savePoleLeash(BlockPos pos, ResourceKey<Level> dimension) {
        this.savedPolePosition = pos;
        this.savedPoleDimension = dimension;
    }

    @Override
    public void clearSavedPoleLeash() {
        this.savedPolePosition = null;
        this.savedPoleDimension = null;
    }

    // Captor persistence

    @Override
    public boolean hasSavedCaptor() {
        return savedCaptorUUID != null;
    }

    @Override
    @Nullable
    public UUID getSavedCaptorUUID() {
        return savedCaptorUUID;
    }

    @Override
    public void saveCaptorUUID(UUID uuid) {
        this.savedCaptorUUID = uuid;
    }

    @Override
    public void clearSavedCaptor() {
        this.savedCaptorUUID = null;
    }

    // NBT serialization

    @Override
    public CompoundTag serializeNBT() {
        CompoundTag root = new CompoundTag();
        CompoundTag regionsTag = new CompoundTag();

        // Track which stacks we've already serialized (identity-based)
        IdentityHashMap<ItemStack, String> serialized = new IdentityHashMap<>();

        for (BodyRegionV2 region : BodyRegionV2.values()) {
            ItemStack stack = regions.get(region);
            if (stack == null || stack.isEmpty()) continue;

            String regionName = region.name();

            if (serialized.containsKey(stack)) {
                // This stack was already serialized under another region.
                // Add this region to the _also list of the primary region.
                String primaryRegion = serialized.get(stack);
                String alsoKey = primaryRegion + NBT_ALSO_SUFFIX;
                ListTag alsoList;
                if (regionsTag.contains(alsoKey, Tag.TAG_LIST)) {
                    alsoList = regionsTag.getList(alsoKey, Tag.TAG_STRING);
                } else {
                    alsoList = new ListTag();
                    regionsTag.put(alsoKey, alsoList);
                }
                alsoList.add(StringTag.valueOf(regionName));
            } else {
                // First time seeing this stack — serialize it
                regionsTag.put(regionName, stack.save(new CompoundTag()));
                serialized.put(stack, regionName);
            }
        }

        root.put(NBT_ROOT_KEY, regionsTag);

        // Pole leash persistence
        if (savedPolePosition != null && savedPoleDimension != null) {
            root.putLong("pole_position", savedPolePosition.asLong());
            root.putString(
                "pole_dimension",
                savedPoleDimension.location().toString()
            );
        }

        // Captor persistence
        if (savedCaptorUUID != null) {
            root.putUUID("captor_uuid", savedCaptorUUID);
        }

        return root;
    }

    @Override
    public void deserializeNBT(CompoundTag tag) {
        clearAll();

        if (!tag.contains(NBT_ROOT_KEY, Tag.TAG_COMPOUND)) return;
        CompoundTag regionsTag = tag.getCompound(NBT_ROOT_KEY);

        // First pass: load primary region stacks
        Set<String> allKeys = regionsTag.getAllKeys();
        Map<String, ItemStack> loadedStacks = new LinkedHashMap<>();

        for (String key : allKeys) {
            if (key.endsWith(NBT_ALSO_SUFFIX)) continue; // Handle in second pass

            BodyRegionV2 region = BodyRegionV2.fromName(key);
            if (region == null) continue;

            CompoundTag stackTag = regionsTag.getCompound(key);
            ItemStack stack = ItemStack.of(stackTag);
            if (stack.isEmpty()) continue;

            regions.put(region, stack);
            loadedStacks.put(key, stack);
        }

        // Second pass: process _also entries to share the same ItemStack reference
        for (String key : allKeys) {
            if (!key.endsWith(NBT_ALSO_SUFFIX)) continue;

            String primaryRegionName = key.substring(
                0,
                key.length() - NBT_ALSO_SUFFIX.length()
            );
            ItemStack primaryStack = loadedStacks.get(primaryRegionName);
            if (primaryStack == null) continue;

            ListTag alsoList = regionsTag.getList(key, Tag.TAG_STRING);
            for (int i = 0; i < alsoList.size(); i++) {
                String alsoRegionName = alsoList.getString(i);
                BodyRegionV2 alsoRegion = BodyRegionV2.fromName(alsoRegionName);
                if (alsoRegion != null) {
                    regions.put(alsoRegion, primaryStack); // Same reference
                }
            }
        }

        // Pole leash persistence
        if (tag.contains("pole_position") && tag.contains("pole_dimension")) {
            try {
                savedPolePosition = BlockPos.of(tag.getLong("pole_position"));
                savedPoleDimension = ResourceKey.create(
                    Registries.DIMENSION,
                    ResourceLocation.parse(tag.getString("pole_dimension"))
                );
            } catch (net.minecraft.ResourceLocationException e) {
                com.tiedup.remake.core.TiedUpMod.LOGGER.warn(
                    "Invalid pole dimension in NBT, clearing saved pole data: {}",
                    e.getMessage()
                );
                savedPolePosition = null;
                savedPoleDimension = null;
            }
        } else {
            savedPolePosition = null;
            savedPoleDimension = null;
        }

        // Captor persistence
        if (tag.hasUUID("captor_uuid")) {
            savedCaptorUUID = tag.getUUID("captor_uuid");
        } else {
            savedCaptorUUID = null;
        }
    }
}
