// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Threading;

namespace Microsoft.NET.StringTools
{
    /// <summary>
    /// Implements the WeakStringCache functionality on modern .NET versions where ConcurrentDictionary is available.
    /// </summary>
    internal sealed partial class WeakStringCache : IDisposable
    {
        private const int WeakHandleMinimumLength = 500;
        private readonly ConcurrentDictionary<int, string> _stringsByHashCode;
        private readonly ConcurrentDictionary<int, StringWeakHandle> _weakHandlesByHashCode;
        private int _count;

        public WeakStringCache()
        {
            _stringsByHashCode = new ConcurrentDictionary<int, string>(Environment.ProcessorCount, _initialCapacity);
            _weakHandlesByHashCode = new ConcurrentDictionary<int, StringWeakHandle>(Environment.ProcessorCount, _initialCapacity);
        }

        /// <summary>
        /// Main entrypoint of this cache. Tries to look up a string that matches the given internable. If it succeeds, returns
        /// the string and sets cacheHit to true. If the string is not found, calls ExpensiveConvertToString on the internable,
        /// adds the resulting string to the cache, and returns it, setting cacheHit to false.
        /// </summary>
        /// <param name="internable">The internable describing the string we're looking for.</param>
        /// <param name="cacheHit">true if match found in cache, false otherwise.</param>
        /// <returns>A string matching the given internable.</returns>
        public string GetOrCreateEntry(ref InternableString internable, out bool cacheHit)
        {
            int hashCode = internable.GetHashCode();

            // We use separate caches for string references and weak handles, as only the latter requires taking a lock
            // on the handle.
            bool usingWeakHandle = internable.Length > WeakHandleMinimumLength;
            return usingWeakHandle ?
                GetStringFromWeakHandle(ref internable, out cacheHit, hashCode) :
                GetString(ref internable, out cacheHit, hashCode);

            string GetString(ref InternableString internable, out bool cacheHit, int hashCode)
            {
                ConcurrentDictionary<int, string> stringsByHashCode = _stringsByHashCode;
                if (stringsByHashCode.TryGetValue(hashCode, out string? result))
                {
                    if (internable.Equals(result))
                    {
                        cacheHit = true;
                        return result;
                    }
                }

                cacheHit = false;
                result = internable.ExpensiveConvertToString();
                stringsByHashCode[hashCode] = result; // Update the cache with the new string.

                return result;
            }

            string GetStringFromWeakHandle(ref InternableString internable, out bool cacheHit, int hashCode)
            {
                string? result;
                ConcurrentDictionary<int, StringWeakHandle> weakHandlesByHashCode = _weakHandlesByHashCode;
                if (weakHandlesByHashCode.TryGetValue(hashCode, out StringWeakHandle? handle))
                {
                    // Lock a weak handle while we're dereferencing it to prevent a race with the Scavenge
                    // method running on another thread and freeing the handle from underneath us.
                    Monitor.Enter(handle);

                    try
                    {
                        result = handle.GetString(ref internable);
                        if (result != null)
                        {
                            cacheHit = true;
                            return result;
                        }

                        // We have the handle but it's not referencing the right string - create the right string and store it in the handle.
                        result = internable.ExpensiveConvertToString();
                        handle.SetString(result);

                        cacheHit = false;
                        return result;
                    }
                    finally
                    {
                        Monitor.Exit(handle);
                    }
                }

                // We don't have the handle in the cache - create the right string, store it in the handle, and add the handle to the cache.
                result = internable.ExpensiveConvertToString();

                handle = new StringWeakHandle();
                handle.SetString(result);
                if (weakHandlesByHashCode.TryAdd(hashCode, handle))
                {
                    // Scavenge only clears out weak handles.
                    Interlocked.Increment(ref _count);
                }

                // Remove unused handles if our heuristic indicates that it would be productive.
                int scavengeThreshold = _scavengeThreshold;
                if (_count >= scavengeThreshold)
                {
                    // Before we start scavenging set _scavengeThreshold to a high value to effectively lock other threads from
                    // running Scavenge at the same time.
                    if (Interlocked.CompareExchange(ref _scavengeThreshold, int.MaxValue, scavengeThreshold) == scavengeThreshold)
                    {
                        try
                        {
                            // Get rid of unused handles.
                            Scavenge();
                        }
                        finally
                        {
                            // And do this again when the number of handles reaches double the current after-scavenge number.
                            _scavengeThreshold = _weakHandlesByHashCode.Count * 2;

                            // This count is not exact, since there can be some Interlocked.Increment(ref _count);
                            // calls happening due to this not being behind a lock.
                            // e.g. code checks if (_stringsByHashCode.TryAdd(hashCode, handle)), we set the _count here and the code increments
                            // however since this is just a threshold to scavenge, it should be fine to be off by few even if that happens.
                            _count = _weakHandlesByHashCode.Count;
                        }
                    }
                }

                cacheHit = false;
                return result;
            }
        }

        /// <summary>
        /// Iterates over the cache and removes unused GC handles, i.e. handles that don't reference live strings.
        /// This is expensive so try to call such that the cost is amortized to O(1) per GetOrCreateEntry() invocation.
        /// </summary>
        public void Scavenge()
        {
            foreach (KeyValuePair<int, StringWeakHandle> entry in _weakHandlesByHashCode)
            {
                // We can safely dereference entry.Value as the caller guarantees that Scavenge runs only on one thread.
                if (!entry.Value.IsUsed && _weakHandlesByHashCode.TryRemove(entry.Key, out StringWeakHandle? removedHandle))
                {
                    lock (removedHandle)
                    {
                        // Note that the removed handle may be different from the one we got from the enumerator so check again
                        // and try to put it back if it's still in use.
                        if (!removedHandle.IsUsed || !_weakHandlesByHashCode.TryAdd(entry.Key, removedHandle))
                        {
                            removedHandle.Free();
                        }
                    }
                }
            }
        }

        /// <summary>
        /// Returns internal debug counters calculated based on the current state of the cache.
        /// </summary>
        public DebugInfo GetDebugInfo()
        {
            return GetDebugInfoImpl();
        }
    }
}