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

using System.Buffers;
using System.Diagnostics.CodeAnalysis;
using System.Text;

namespace Microsoft.AspNetCore.SignalR.Internal;

/// <summary>
/// A small dictionary optimized for utf8 string lookup via spans. Adapted from https://github.com/dotnet/runtime/blob/4ed596ef63e60ce54cfb41d55928f0fe45f65cf3/src/libraries/System.Linq.Parallel/src/System/Linq/Parallel/Utils/HashLookup.cs.
/// </summary>
internal sealed class Utf8HashLookup
{
    private int[] _buckets;
    private int[] _caseSensitiveBuckets;
    private Slot[] _slots;
    private int _count;

    private const int HashCodeMask = 0x7fffffff;

    internal Utf8HashLookup()
    {
        _buckets = new int[7];
        _caseSensitiveBuckets = new int[7];
        _slots = new Slot[7];
    }

    internal void Add(string value)
    {
        if (_count == _slots.Length)
        {
            Resize();
        }

        int slotIndex = _count;
        _count++;

        var encodedValue = Encoding.UTF8.GetBytes(value);
        var hashCode = GetHashCode(value.AsSpan());
        var caseSensitiveHashCode = GetCaseSensitiveHashCode(encodedValue);
        int bucketIndex = hashCode % _buckets.Length;
        int caseSensitiveBucketIndex = caseSensitiveHashCode % _caseSensitiveBuckets.Length;

        _slots[slotIndex].hashCode = hashCode;
        _slots[slotIndex].caseSensitiveHashCode = caseSensitiveHashCode;

        _slots[slotIndex].value = value;
        _slots[slotIndex].encodedValue = encodedValue;

        _slots[slotIndex].next = _buckets[bucketIndex] - 1;
        _slots[slotIndex].caseSensitiveNext = _caseSensitiveBuckets[caseSensitiveBucketIndex] - 1;

        _buckets[bucketIndex] = slotIndex + 1;
        _caseSensitiveBuckets[caseSensitiveBucketIndex] = slotIndex + 1;
    }

    internal bool TryGetValue(ReadOnlySpan<byte> encodedValue, [MaybeNullWhen(false), AllowNull] out string value)
    {
        var caseSensitiveHashCode = GetCaseSensitiveHashCode(encodedValue);

        for (var i = _caseSensitiveBuckets[caseSensitiveHashCode % _caseSensitiveBuckets.Length] - 1; i >= 0; i = _slots[i].caseSensitiveNext)
        {
            if (_slots[i].caseSensitiveHashCode == caseSensitiveHashCode && encodedValue.SequenceEqual(_slots[i].encodedValue.AsSpan()))
            {
                value = _slots[i].value;
                return true;
            }
        }

        // If we cannot find a case-sensitive match, we transcode the encodedValue to a stackalloced UTF16 string
        // and do an OrdinalIgnoreCase comparison.
        return TryGetValueSlow(encodedValue, out value);
    }

    private bool TryGetValueSlow(ReadOnlySpan<byte> encodedValue, [MaybeNullWhen(false), AllowNull] out string value)
    {
        const int StackAllocThreshold = 128;

        char[]? pooled = null;
        var count = Encoding.UTF8.GetCharCount(encodedValue);
        var chars = count <= StackAllocThreshold ?
            stackalloc char[StackAllocThreshold] :
            (pooled = ArrayPool<char>.Shared.Rent(count));
        var encoded = Encoding.UTF8.GetChars(encodedValue, chars);
        var hasValue = TryGetValueFromChars(chars[..encoded], out value);
        if (pooled is not null)
        {
            ArrayPool<char>.Shared.Return(pooled);
        }

        return hasValue;
    }

    private bool TryGetValueFromChars(ReadOnlySpan<char> key, [MaybeNullWhen(false), AllowNull] out string value)
    {
        var hashCode = GetHashCode(key);

        for (var i = _buckets[hashCode % _buckets.Length] - 1; i >= 0; i = _slots[i].next)
        {
            if (_slots[i].hashCode == hashCode && key.Equals(_slots[i].value, StringComparison.OrdinalIgnoreCase))
            {
                value = _slots[i].value;
                return true;
            }
        }

        value = null;
        return false;
    }

    private static int GetHashCode(ReadOnlySpan<char> value) =>
        HashCodeMask & string.GetHashCode(value, StringComparison.OrdinalIgnoreCase);

    private static int GetCaseSensitiveHashCode(ReadOnlySpan<byte> encodedValue)
    {
        var hashCode = new HashCode();
        hashCode.AddBytes(encodedValue);
        return HashCodeMask & hashCode.ToHashCode();
    }

    private void Resize()
    {
        var newSize = checked(_count * 2 + 1);
        var newSlots = new Slot[newSize];

        var newBuckets = new int[newSize];
        var newCaseSensitiveBuckets = new int[newSize];

        Array.Copy(_slots, newSlots, _count);

        for (int i = 0; i < _count; i++)
        {
            int bucket = newSlots[i].hashCode % newSize;
            newSlots[i].next = newBuckets[bucket] - 1;
            newBuckets[bucket] = i + 1;

            int caseSensitiveBucket = newSlots[i].caseSensitiveHashCode % newSize;
            newSlots[i].caseSensitiveNext = newCaseSensitiveBuckets[caseSensitiveBucket] - 1;
            newCaseSensitiveBuckets[caseSensitiveBucket] = i + 1;
        }

        _slots = newSlots;

        _buckets = newBuckets;
        _caseSensitiveBuckets = newCaseSensitiveBuckets;
    }

    private struct Slot
    {
        internal int hashCode;
        internal int caseSensitiveHashCode;

        internal string value;
        internal byte[] encodedValue;

        internal int next;
        internal int caseSensitiveNext;
    }
}