megakoop/Game/Scripts/Networking/NetworkIdRegistry.cs

using System.Collections.Generic;
using UnityEngine;

namespace MegaKoop.Game.Networking
{
    /// <summary>
    /// Centralized registry for managing Network IDs across all clients.
    /// Provides collision detection, Steam ID mapping, and ID range reservation.
    /// </summary>
    public class NetworkIdRegistry : MonoBehaviour
    {
        private static NetworkIdRegistry instance;

        private readonly Dictionary<int, NetworkIdentity> registry = new();
        private readonly Dictionary<ulong, int> steamIdToNetworkId = new();
        private readonly HashSet<int> reservedIds = new();

        public static NetworkIdRegistry Instance
        {
            get
            {
                if (instance == null)
                {
                    GameObject go = new GameObject("NetworkIdRegistry");
                    instance = go.AddComponent<NetworkIdRegistry>();
                    DontDestroyOnLoad(go);
                }
                return instance;
            }
        }

        private void Awake()
        {
            if (instance != null && instance != this)
            {
                Debug.LogWarning("[NetworkIdRegistry] Duplicate instance detected. Destroying.");
                Destroy(gameObject);
                return;
            }

            instance = this;
            DontDestroyOnLoad(gameObject);
        }

        /// <summary>
        /// Attempts to register a NetworkIdentity with its assigned ID.
        /// </summary>
        /// <param name="identity">The NetworkIdentity to register</param>
        /// <returns>True if registration succeeded, false if collision detected</returns>
        public static bool TryRegister(NetworkIdentity identity)
        {
            if (identity == null)
            {
                Debug.LogError("[NetworkIdRegistry] Cannot register null NetworkIdentity.");
                return false;
            }

            int networkId = identity.NetworkId;

            if (networkId == 0)
            {
                Debug.LogWarning($"[NetworkIdRegistry] NetworkIdentity '{identity.name}' has ID 0 and cannot be registered.");
                return false;
            }

            var inst = Instance;

            // Check for collision
            if (inst.registry.TryGetValue(networkId, out NetworkIdentity existing))
            {
                if (existing == identity)
                {
                    // Already registered
                    return true;
                }

                if (existing != null)
                {
                    Debug.LogError($"[NetworkIdRegistry] ID collision detected! ID {networkId} is already assigned to '{existing.name}'. " +
                                   $"Cannot register '{identity.name}'. This will cause network synchronization issues.");
                    return false;
                }
            }

            // Check if ID is in reserved range
            if (inst.reservedIds.Contains(networkId))
            {
                Debug.LogWarning($"[NetworkIdRegistry] ID {networkId} is in a reserved range but being registered for '{identity.name}'.");
            }

            inst.registry[networkId] = identity;
            Debug.Log($"[NetworkIdRegistry] Registered '{identity.name}' with Network ID {networkId}");
            return true;
        }

        /// <summary>
        /// Unregisters a NetworkIdentity by its ID.
        /// </summary>
        /// <param name="networkId">The network ID to unregister</param>
        public static void Unregister(int networkId)
        {
            if (instance == null)
            {
                return;
            }

            if (instance.registry.Remove(networkId))
            {
                Debug.Log($"[NetworkIdRegistry] Unregistered Network ID {networkId}");
            }

            SteamCharacterStateCache.RemoveState(networkId);

            // Also remove from Steam ID mapping if present
            ulong steamIdToRemove = 0;
            foreach (var kvp in instance.steamIdToNetworkId)
            {
                if (kvp.Value == networkId)
                {
                    steamIdToRemove = kvp.Key;
                    break;
                }
            }

            if (steamIdToRemove != 0)
            {
                instance.steamIdToNetworkId.Remove(steamIdToRemove);
            }
        }

        /// <summary>
        /// Retrieves a NetworkIdentity by its ID.
        /// </summary>
        /// <param name="networkId">The network ID to look up</param>
        /// <returns>The NetworkIdentity if found, null otherwise</returns>
        public static NetworkIdentity GetById(int networkId)
        {
            if (instance == null)
            {
                return null;
            }

            instance.registry.TryGetValue(networkId, out NetworkIdentity identity);
            return identity;
        }

        /// <summary>
        /// Maps a Steam ID to a Network ID for player tracking.
        /// </summary>
        /// <param name="steamId">The Steam ID of the player</param>
        /// <param name="networkId">The Network ID assigned to the player</param>
        public static void MapSteamIdToNetworkId(ulong steamId, int networkId)
        {
            if (steamId == 0)
            {
                Debug.LogWarning("[NetworkIdRegistry] Cannot map Steam ID 0.");
                return;
            }

            var inst = Instance;

            if (inst.steamIdToNetworkId.TryGetValue(steamId, out int existingId))
            {
                if (existingId != networkId)
                {
                    Debug.LogWarning($"[NetworkIdRegistry] Steam ID {steamId} was mapped to Network ID {existingId}, now remapping to {networkId}");
                }
            }

            inst.steamIdToNetworkId[steamId] = networkId;
            Debug.Log($"[NetworkIdRegistry] Mapped Steam ID {steamId} to Network ID {networkId}");
        }

        /// <summary>
        /// Gets the Network ID associated with a Steam ID.
        /// </summary>
        /// <param name="steamId">The Steam ID to look up</param>
        /// <returns>The Network ID if found, 0 otherwise</returns>
        public static int GetNetworkIdForSteamId(ulong steamId)
        {
            if (instance == null || steamId == 0)
            {
                return 0;
            }

            instance.steamIdToNetworkId.TryGetValue(steamId, out int networkId);
            return networkId;
        }

        /// <summary>
        /// Reserves a range of IDs to prevent automatic assignment.
        /// </summary>
        /// <param name="startId">The starting ID of the range (inclusive)</param>
        /// <param name="count">The number of IDs to reserve</param>
        public static void ReserveIdRange(int startId, int count)
        {
            if (startId <= 0 || count <= 0)
            {
                Debug.LogWarning($"[NetworkIdRegistry] Invalid ID range: start={startId}, count={count}");
                return;
            }

            var inst = Instance;

            for (int i = 0; i < count; i++)
            {
                inst.reservedIds.Add(startId + i);
            }

            Debug.Log($"[NetworkIdRegistry] Reserved ID range {startId} to {startId + count - 1} ({count} IDs)");
        }

        /// <summary>
        /// Checks if an ID is in a reserved range.
        /// </summary>
        /// <param name="networkId">The ID to check</param>
        /// <returns>True if the ID is reserved, false otherwise</returns>
        public static bool IsIdReserved(int networkId)
        {
            if (instance == null)
            {
                return false;
            }

            return instance.reservedIds.Contains(networkId);
        }

        /// <summary>
        /// Clears all registrations and mappings. Use when transitioning between scenes or sessions.
        /// </summary>
        public static void ClearRegistry()
        {
            if (instance == null)
            {
                return;
            }

            int registryCount = instance.registry.Count;
            int mappingCount = instance.steamIdToNetworkId.Count;
            int reservedCount = instance.reservedIds.Count;

            instance.registry.Clear();
            instance.steamIdToNetworkId.Clear();
            instance.reservedIds.Clear();

            Debug.Log($"[NetworkIdRegistry] Registry cleared. Removed {registryCount} identities, {mappingCount} Steam ID mappings, and {reservedCount} reserved IDs.");
        }

        /// <summary>
        /// Gets the total number of registered identities.
        /// </summary>
        public static int GetRegisteredCount()
        {
            return instance?.registry.Count ?? 0;
        }

        /// <summary>
        /// Checks if a specific Network ID is currently registered.
        /// </summary>
        public static bool IsIdRegistered(int networkId)
        {
            return instance != null && instance.registry.ContainsKey(networkId);
        }

        private void OnDestroy()
        {
            if (instance == this)
            {
                instance = null;
            }
        }
    }
}