TiedUp-/src/main/java/com/tiedup/remake/state/CollarRegistry.java

package com.tiedup.remake.state;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.ListTag;
import net.minecraft.nbt.Tag;
import net.minecraft.server.MinecraftServer;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.server.level.ServerPlayer;
import net.minecraft.world.entity.Entity;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.level.saveddata.SavedData;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

/**
 * Global registry for collar ownership relationships.
 *
 * This registry tracks which entities are wearing collars and who owns them.
 * It persists across server restarts and provides efficient lookups in both directions:
 * - Owner UUID → Set of collar-wearer UUIDs (slaves)
 * - Wearer UUID → Set of owner UUIDs (masters)
 *
 * Terminology:
 * - "Slave" = Entity wearing a collar owned by a player (passive ownership)
 * - "Captive" = Entity attached by leash (active physical control) - managed by PlayerCaptiveManager
 *
 */
public class CollarRegistry extends SavedData {

    private static final String DATA_NAME = "tiedup_collar_registry";

    // Owner UUID → Set of wearer UUIDs (a master can own multiple slaves)
    private final Map<UUID, Set<UUID>> ownerToWearers =
        new ConcurrentHashMap<>();

    // Wearer UUID → Set of owner UUIDs (a collar can have multiple owners)
    private final Map<UUID, Set<UUID>> wearerToOwners =
        new ConcurrentHashMap<>();

    // ==================== STATIC ACCESS ====================

    /**
     * Get the CollarRegistry for a server level.
     * Creates a new registry if one doesn't exist.
     */
    public static CollarRegistry get(ServerLevel level) {
        return level
            .getDataStorage()
            .computeIfAbsent(
                CollarRegistry::load,
                CollarRegistry::new,
                DATA_NAME
            );
    }

    /**
     * Get the CollarRegistry from a MinecraftServer.
     * Uses the overworld as the storage dimension.
     */
    public static CollarRegistry get(MinecraftServer server) {
        ServerLevel overworld = server.overworld();
        return get(overworld);
    }

    /**
     * Convenience method to get registry from a ServerPlayer.
     */
    @Nullable
    public static CollarRegistry get(ServerPlayer player) {
        if (player.getServer() == null) return null;
        return get(player.getServer());
    }

    // ==================== REGISTRATION ====================

    /**
     * Register a collar relationship: owner now owns the collar on wearer.
     *
     * @param wearerUUID UUID of the entity wearing the collar
     * @param ownerUUID UUID of the collar's owner
     */
    public void registerCollar(UUID wearerUUID, UUID ownerUUID) {
        // Add to owner → wearers map
        ownerToWearers
            .computeIfAbsent(ownerUUID, k -> ConcurrentHashMap.newKeySet())
            .add(wearerUUID);

        // Add to wearer → owners map
        wearerToOwners
            .computeIfAbsent(wearerUUID, k -> ConcurrentHashMap.newKeySet())
            .add(ownerUUID);

        setDirty();
    }

    /**
     * Register a collar with multiple owners at once.
     *
     * @param wearerUUID UUID of the entity wearing the collar
     * @param ownerUUIDs Set of owner UUIDs
     */
    public void registerCollar(UUID wearerUUID, Set<UUID> ownerUUIDs) {
        for (UUID ownerUUID : ownerUUIDs) {
            registerCollar(wearerUUID, ownerUUID);
        }
    }

    /**
     * Unregister a specific owner from a collar wearer.
     *
     * @param wearerUUID UUID of the entity wearing the collar
     * @param ownerUUID UUID of the owner to remove
     */
    private void unregisterOwner(UUID wearerUUID, UUID ownerUUID) {
        // Remove from owner → wearers map using atomic operation
        ownerToWearers.computeIfPresent(ownerUUID, (key, wearers) -> {
            wearers.remove(wearerUUID);
            return wearers.isEmpty() ? null : wearers;
        });

        // Remove from wearer → owners map using atomic operation
        wearerToOwners.computeIfPresent(wearerUUID, (key, owners) -> {
            owners.remove(ownerUUID);
            return owners.isEmpty() ? null : owners;
        });

        setDirty();
    }

    /**
     * Completely unregister a collar wearer (removes all owner relationships).
     * Called when a collar is removed from an entity.
     *
     * @param wearerUUID UUID of the entity whose collar was removed
     */
    public void unregisterWearer(UUID wearerUUID) {
        Set<UUID> owners = wearerToOwners.remove(wearerUUID);
        if (owners != null) {
            for (UUID ownerUUID : owners) {
                // Use atomic operation for thread-safe removal
                ownerToWearers.computeIfPresent(ownerUUID, (key, wearers) -> {
                    wearers.remove(wearerUUID);
                    return wearers.isEmpty() ? null : wearers;
                });
            }
            setDirty();
        }
    }

    /**
     * Update a wearer's owners completely (replaces all existing owners).
     * Useful when collar NBT is the source of truth.
     *
     * @param wearerUUID UUID of the entity wearing the collar
     * @param newOwnerUUIDs New set of owner UUIDs
     */
    public void updateWearerOwners(UUID wearerUUID, Set<UUID> newOwnerUUIDs) {
        // Get current owners
        Set<UUID> currentOwners = wearerToOwners.get(wearerUUID);
        if (currentOwners == null) {
            currentOwners = Collections.emptySet();
        }

        // Find owners to remove
        Set<UUID> toRemove = new HashSet<>(currentOwners);
        toRemove.removeAll(newOwnerUUIDs);

        // Find owners to add
        Set<UUID> toAdd = new HashSet<>(newOwnerUUIDs);
        toAdd.removeAll(currentOwners);

        // Apply changes
        for (UUID ownerUUID : toRemove) {
            unregisterOwner(wearerUUID, ownerUUID);
        }
        for (UUID ownerUUID : toAdd) {
            registerCollar(wearerUUID, ownerUUID);
        }
    }

    // ==================== QUERIES ====================

    /**
     * Get all slaves (collar wearers) owned by a specific owner.
     *
     * @param ownerUUID UUID of the owner
     * @return Unmodifiable set of wearer UUIDs (never null)
     */
    public Set<UUID> getSlaves(UUID ownerUUID) {
        Set<UUID> wearers = ownerToWearers.get(ownerUUID);
        if (wearers == null) return Collections.emptySet();
        return Collections.unmodifiableSet(new HashSet<>(wearers));
    }

    /**
     * Get all owners of a specific collar wearer.
     *
     * @param wearerUUID UUID of the wearer
     * @return Unmodifiable set of owner UUIDs (never null)
     */
    public Set<UUID> getOwners(UUID wearerUUID) {
        Set<UUID> owners = wearerToOwners.get(wearerUUID);
        if (owners == null) return Collections.emptySet();
        return Collections.unmodifiableSet(new HashSet<>(owners));
    }

    /**
     * Check if an owner has any slaves.
     */
    public boolean hasSlaves(UUID ownerUUID) {
        Set<UUID> wearers = ownerToWearers.get(ownerUUID);
        return wearers != null && !wearers.isEmpty();
    }

    /**
     * Check if a wearer has any owners.
     */
    public boolean hasOwners(UUID wearerUUID) {
        Set<UUID> owners = wearerToOwners.get(wearerUUID);
        return owners != null && !owners.isEmpty();
    }

    /**
     * Check if a specific owner owns a specific wearer.
     */
    public boolean isOwner(UUID ownerUUID, UUID wearerUUID) {
        Set<UUID> wearers = ownerToWearers.get(ownerUUID);
        return wearers != null && wearers.contains(wearerUUID);
    }

    /**
     * Get the count of slaves for an owner.
     */
    public int getSlaveCount(UUID ownerUUID) {
        Set<UUID> wearers = ownerToWearers.get(ownerUUID);
        return wearers == null ? 0 : wearers.size();
    }

    /**
     * Get all registered wearers (for admin/debug purposes).
     */
    public Set<UUID> getAllWearers() {
        return Collections.unmodifiableSet(
            new HashSet<>(wearerToOwners.keySet())
        );
    }

    /**
     * Get all registered owners (for admin/debug purposes).
     */
    public Set<UUID> getAllOwners() {
        return Collections.unmodifiableSet(
            new HashSet<>(ownerToWearers.keySet())
        );
    }

    // ==================== ENTITY RESOLUTION ====================

    /**
     * Find all slave entities for an owner that are currently loaded.
     * This is useful for GUI and proximity-based actions.
     *
     * @param owner The owner player
     * @return List of currently loaded slave entities
     */
    public List<LivingEntity> findLoadedSlaves(ServerPlayer owner) {
        List<LivingEntity> loadedSlaves = new ArrayList<>();
        Set<UUID> slaveUUIDs = getSlaves(owner.getUUID());

        MinecraftServer server = owner.getServer();
        if (server == null) return loadedSlaves;

        for (UUID slaveUUID : slaveUUIDs) {
            Entity entity = findEntityByUUID(server, slaveUUID);
            if (entity instanceof LivingEntity living) {
                loadedSlaves.add(living);
            }
        }

        return loadedSlaves;
    }

    /**
     * Find an entity by UUID across all dimensions.
     * Optimized: checks player list first (O(1)) before searching dimensions.
     */
    @Nullable
    private Entity findEntityByUUID(MinecraftServer server, UUID uuid) {
        // Check player first - O(1) lookup
        net.minecraft.server.level.ServerPlayer player = server
            .getPlayerList()
            .getPlayer(uuid);
        if (player != null) {
            return player;
        }

        // Fallback: search dimensions for NPCs
        for (ServerLevel level : server.getAllLevels()) {
            Entity entity = level.getEntity(uuid);
            if (entity != null) {
                return entity;
            }
        }
        return null;
    }

    // ==================== PERSISTENCE ====================

    @Override
    public @NotNull CompoundTag save(@NotNull CompoundTag tag) {
        ListTag registryList = new ListTag();

        for (Map.Entry<UUID, Set<UUID>> entry : wearerToOwners.entrySet()) {
            CompoundTag wearerTag = new CompoundTag();
            wearerTag.putUUID("wearer", entry.getKey());

            ListTag ownersTag = new ListTag();
            for (UUID ownerUUID : entry.getValue()) {
                CompoundTag ownerTag = new CompoundTag();
                ownerTag.putUUID("uuid", ownerUUID);
                ownersTag.add(ownerTag);
            }
            wearerTag.put("owners", ownersTag);
            registryList.add(wearerTag);
        }

        tag.put("collar_registry", registryList);
        return tag;
    }

    public static CollarRegistry load(CompoundTag tag) {
        CollarRegistry registry = new CollarRegistry();

        ListTag registryList = tag.getList("collar_registry", Tag.TAG_COMPOUND);
        for (int i = 0; i < registryList.size(); i++) {
            CompoundTag wearerTag = registryList.getCompound(i);
            UUID wearerUUID = wearerTag.getUUID("wearer");

            ListTag ownersTag = wearerTag.getList("owners", Tag.TAG_COMPOUND);
            for (int j = 0; j < ownersTag.size(); j++) {
                CompoundTag ownerTag = ownersTag.getCompound(j);
                UUID ownerUUID = ownerTag.getUUID("uuid");

                // Register relationship (without marking dirty - we're loading)
                registry.ownerToWearers
                    .computeIfAbsent(ownerUUID, k ->
                        ConcurrentHashMap.newKeySet()
                    )
                    .add(wearerUUID);
                registry.wearerToOwners
                    .computeIfAbsent(wearerUUID, k ->
                        ConcurrentHashMap.newKeySet()
                    )
                    .add(ownerUUID);
            }
        }

        return registry;
    }

    // ==================== DEBUG ====================

    /**
     * Get a debug string representation of the registry.
     */
    public String toDebugString() {
        StringBuilder sb = new StringBuilder();
        sb.append("CollarRegistry:\n");
        sb.append("  Owners: ").append(ownerToWearers.size()).append("\n");
        sb.append("  Wearers: ").append(wearerToOwners.size()).append("\n");

        for (Map.Entry<UUID, Set<UUID>> entry : ownerToWearers.entrySet()) {
            sb
                .append("  Owner ")
                .append(entry.getKey().toString().substring(0, 8))
                .append("... → ")
                .append(entry.getValue().size())
                .append(" slaves\n");
        }

        return sb.toString();
    }
}