using Microsoft.Toolkit.HighPerformance;
using SpriteMaster.Extensions;
using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading;

namespace SpriteMaster.Types;

internal class TypedMemoryCache<TKey, TValue> : AbstractTypedMemoryCache<TKey, TValue> where TKey : notnull where TValue : class {
    [StructLayout(LayoutKind.Auto)]
    private readonly struct CacheEntry {
        internal readonly TValue Value;
        internal readonly ConcurrentLinkedListSlim<TKey>.NodeRef Node;
        internal readonly long Size;

        internal readonly bool IsValid => Node.IsValid;

        [MethodImpl(Runtime.MethodImpl.Inline)]
        internal CacheEntry(TValue value, ConcurrentLinkedListSlim<TKey>.NodeRef node, long size) {
            Value = value;
            Node = node;
            Size = size;
        }

        [MethodImpl(Runtime.MethodImpl.Inline)]
        internal readonly void UpdateValue(TValue value, long size) {
            Unsafe.AsRef(in Value) = value;
            Unsafe.AsRef(in Size) = size;
        }
    }

    private readonly Dictionary<TKey, CacheEntry> Cache = new();
    private readonly ConcurrentLinkedListSlim<TKey> RecentAccessList = new();
    private readonly SharedLock CacheLock = new();

    private readonly CancellationTokenSource ThreadCancelSource = new();
    private readonly Condition TrimEvent = new();
    private readonly Thread TrimThread;

    private void CacheTrimLoop(in CancellationToken cancellationToken) {
        while (!cancellationToken.IsCancellationRequested) {
            TrimEvent.Wait();

            if (cancellationToken.IsCancellationRequested) {
                return;
            }

            TrimToSize(MaxSizeHysteresis);
        }
    }

    internal TypedMemoryCache(string name, long? maxSize, RemovalCallbackDelegate? removalAction = null) :
        base(name, maxSize ?? long.MaxValue, removalAction) {
        TrimThread = ThreadExt.Run(() => CacheTrimLoop(ThreadCancelSource.Token), background: true, name: $"Cache '{name}' Trim Thread");
    }

    internal override int Count => Cache.Count;

    [MethodImpl(Runtime.MethodImpl.Inline)]
    internal override TValue? Get(TKey key) {
        using (CacheLock.Read) {
            var entry = Cache.GetValueOrDefault(key);
            if (entry.IsValid) {
                RecentAccessList.MoveToFront(entry.Node);
            }

            return default;
        }
    }

    [MethodImpl(Runtime.MethodImpl.Inline)]
    internal override bool TryGet(TKey key, [NotNullWhen(true)] out TValue? value) {
        using (CacheLock.Read) {
            if (Cache.TryGetValue(key, out var entry)) {
                RecentAccessList.MoveToFront(entry.Node);
                value = entry.Value;
                return true;
            }
        }

        value = default;
        return false;
    }

    internal override TValue Set(TKey key, TValue value) {
        TValue? original = null;
        long size = GetSizeBytes(value);
        long sizeDelta;

        using (CacheLock.Write) {
            if (Cache.TryGetValue(key, out var entry)) {
                RecentAccessList.MoveToFront(entry.Node);
                original = entry.Value;
                if (ReferenceEquals(original, value)) {
                    return original;
                }

                long originalSize = entry.Size;
                sizeDelta = size - originalSize;
                entry.UpdateValue(value, size);
                Cache[key] = entry;
            }
            else {
                Cache.Add(key, new(value, RecentAccessList.AddFront(key), size));
                sizeDelta = size;
            }

            Interlocked.Add(ref CurrentSize, sizeDelta);
        }

        if (sizeDelta != 0L) {
            if (Interlocked.Read(ref CurrentSize) > MaxSize) {
                TrimEvent.Set();
            }
        }

        if (original is not null) {
            OnEntryRemoved(key, original, EvictionReason.Replaced);
        }

        return value;
    }

    [MethodImpl(Runtime.MethodImpl.Inline)]
    internal override TValue? Update(TKey key, TValue value) {
        TValue? original = null;
        long size = GetSizeBytes(value);
        long sizeDelta;

        using (CacheLock.ReadWrite) {
            var entry = Cache.GetValueOrDefault(key, default);
            using (CacheLock.Write) {
                if (entry.IsValid) {
                    RecentAccessList.MoveToFront(entry.Node);
                    original = entry.Value;
                    if (ReferenceEquals(original, value)) {
                        return original;
                    }

                    long originalSize = entry.Size;
                    sizeDelta = size - originalSize;
                    entry.UpdateValue(value, size);
                    Cache[key] = entry;
                }
                else {
                    Cache.Add(key, new(value, RecentAccessList.AddFront(key), size));
                    sizeDelta = size;
                }

                Interlocked.Add(ref CurrentSize, sizeDelta);
            }
        }

        if (sizeDelta != 0L) {
            if (Interlocked.Read(ref CurrentSize) > MaxSize) {
                TrimEvent.Set();
            }
        }

        if (original is not null) {
            OnEntryRemoved(key, original, EvictionReason.Replaced);
        }

        return original;
    }

    [MethodImpl(Runtime.MethodImpl.Inline)]
    internal override TValue? Remove(TKey key) {
        TValue? result = null;

        using (CacheLock.Write) {
            if (Cache.Remove(key, out var entry)) {
                result = entry.Value;
                Interlocked.Add(ref CurrentSize, -entry.Size);
                RecentAccessList.Release(entry.Node);
            }
        }

        if (result is not null) {
            OnEntryRemoved(key, result, EvictionReason.Removed);
        }

        return result;
    }

    [MethodImpl(Runtime.MethodImpl.Inline)]
    internal override void Trim(int count) {
        count.AssertPositiveOrZero();

        if (count == 0) {
            return;
        }

        var trimArray = GC.AllocateUninitializedArray<KeyValuePair<TKey, TValue>>(count);

        using (CacheLock.Write) {
            if (count > Cache.Count) {
                count = Cache.Count;
            }

            if (count == 0) {
                return;
            }

            for (int i = 0; i < count; i++) {
                var removeKey = RecentAccessList.RemoveLast();
                bool result = Cache.Remove(removeKey, out var entry);
                Interlocked.Add(ref CurrentSize, -entry.Size);
                trimArray[i] = new(removeKey, entry.Value);
                result.AssertTrue();
            }
        }

        ReportTrimmed(trimArray.AsReadOnlySpan(0, count));
    }

    [MethodImpl(Runtime.MethodImpl.Inline)]
    internal override void TrimTo(int count) {
        count.AssertPositiveOrZero();

        Trim(Cache.Count - count);
    }

    private void TrimToSize(long size) {
        CacheLock.IsWriteLock.AssertFalse();
        size.AssertPositiveOrZero();

        if (size == 0L) {
            Clear();
        }
        else {
            var trimmed = new List<KeyValuePair<TKey, TValue>>();

            using (CacheLock.Write) {
                long currentSize = Interlocked.Read(ref CurrentSize);

                while (currentSize >= size) {
                    var removeKey = RecentAccessList.RemoveLast();
                    bool result = Cache.Remove(removeKey, out var entry);
                    currentSize -= entry.Size;
                    trimmed.Add(new(removeKey, entry.Value));
                    result.AssertTrue();
                }

                CurrentSize = currentSize;
            }

            ReportTrimmed(trimmed.AsSpan());
        }
    }

    private void ReportTrimmed(ReadOnlySpan<KeyValuePair<TKey, TValue>> trimmed) {
        CacheLock.IsWriteLock.AssertFalse();

        foreach (var pair in trimmed) {
            OnEntryRemoved(pair.Key, pair.Value, EvictionReason.Expired);
        }
    }

    [MethodImpl(Runtime.MethodImpl.Inline)]
    internal override void Clear() {
        int removedCount;
        var removedPairs = GC.AllocateUninitializedArray<KeyValuePair<TKey, TValue>>(Count);

        using (CacheLock.Write) {
            if (Count > removedPairs.Length) {
                removedPairs = GC.AllocateUninitializedArray<KeyValuePair<TKey, TValue>>(Count);
            }

            removedCount = Count;

            int removedIndex = 0;
            foreach (var (key, entry) in Cache) {
                removedPairs[removedIndex++] = new(key, entry.Value);
            }

            CurrentSize = 0;
            Cache.Clear();
            RecentAccessList.Clear();
        }

        ReportTrimmed(removedPairs.AsReadOnlySpan(0, removedCount));
    }

    public override void Dispose() {
        ThreadCancelSource.Cancel();
        TrimEvent.Set();

        using (CacheLock.Write) {
            Clear();
        }
        CacheLock.Dispose();

        TrimThread.Join();

        ThreadCancelSource.Dispose();
        TrimEvent.Dispose();
    }
}