using System.Collections.Generic;
using UnityEngine.AddressableAssets;
using UnityEngine.ResourceManagement.AsyncOperations;

namespace Darkmatter.Services.Assets
{
    internal sealed class AddressableLoadHandleTracker
    {
        private readonly Dictionary<string, Stack<TrackedHandleLease>> _leasesByReleaseKey =
            new(System.StringComparer.Ordinal);

        public void Track(AsyncOperationHandle handle, IEnumerable<string> releaseKeys)
        {
            if (!handle.IsValid() || releaseKeys == null)
            {
                return;
            }

            var uniqueReleaseKeys = new List<string>();
            var seenReleaseKeys = new HashSet<string>(System.StringComparer.Ordinal);

            foreach (var releaseKey in releaseKeys)
            {
                var normalizedReleaseKey = Normalize(releaseKey);
                if (string.IsNullOrWhiteSpace(normalizedReleaseKey) || !seenReleaseKeys.Add(normalizedReleaseKey))
                {
                    continue;
                }

                uniqueReleaseKeys.Add(normalizedReleaseKey);
            }

            if (uniqueReleaseKeys.Count == 0)
            {
                return;
            }

            var lease = new TrackedHandleLease(handle, uniqueReleaseKeys.Count);
            for (var i = 0; i < uniqueReleaseKeys.Count; i++)
            {
                var releaseKey = uniqueReleaseKeys[i];
                if (!_leasesByReleaseKey.TryGetValue(releaseKey, out var leases))
                {
                    leases = new Stack<TrackedHandleLease>();
                    _leasesByReleaseKey[releaseKey] = leases;
                }

                leases.Push(lease);
            }
        }

        public bool TryRelease(string releaseKey)
        {
            var normalizedReleaseKey = Normalize(releaseKey);
            if (string.IsNullOrWhiteSpace(normalizedReleaseKey) ||
                !_leasesByReleaseKey.TryGetValue(normalizedReleaseKey, out var leases))
            {
                return false;
            }

            while (leases.Count > 0)
            {
                var lease = leases.Pop();
                if (lease.TryConsume())
                {
                    if (leases.Count == 0)
                    {
                        _leasesByReleaseKey.Remove(normalizedReleaseKey);
                    }

                    return true;
                }
            }

            _leasesByReleaseKey.Remove(normalizedReleaseKey);
            return false;
        }

        private static string Normalize(string releaseKey)
        {
            return releaseKey?.Trim() ?? string.Empty;
        }

        private sealed class TrackedHandleLease
        {
            private readonly AsyncOperationHandle _handle;
            private int _remainingReleaseCount;
            private bool _released;

            public TrackedHandleLease(AsyncOperationHandle handle, int remainingReleaseCount)
            {
                _handle = handle;
                _remainingReleaseCount = remainingReleaseCount;
            }

            public bool TryConsume()
            {
                if (_released || _remainingReleaseCount <= 0)
                {
                    return false;
                }

                _remainingReleaseCount--;
                if (_remainingReleaseCount == 0)
                {
                    if (_handle.IsValid())
                    {
                        Addressables.Release(_handle);
                    }

                    _released = true;
                }

                return true;
            }
        }
    }
}
