XFESystem

using System.Runtime;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;

namespace XFESystem.Collections;

/// <summary>The pinned interpreter implements this interface after its private
/// storage patch. Each engine owns its scope directly; there is no ambient thread state.</summary>
public interface IWeakCollectionOwner { WeakCollectionScope WeakCollections { get; } }

public sealed record WeakCollectionLifetime(int Tables, int Entries, long Created, long Disposed,
    long TableHandlesCreated, long TableHandlesDisposed);

/// <summary>Explicit lifetime for weak collections on runtimes without finalization.
/// The registry holds storage, never a strong table reference. A table's existing
/// key/value semantics therefore cannot be strengthened by registry membership.</summary>
public sealed class WeakCollectionScope
{
    private sealed class Registration(object owner, IWeakStorage storage)
    {
        internal GCHandle Owner = GCHandle.Alloc(owner, GCHandleType.Weak);
        internal readonly IWeakStorage Storage = storage;
    }
    private readonly List<Registration> registrations = [];
    private long created, disposed, tableCreated, tableDisposed;
    private bool released;
    internal void EnsureActive() => ObjectDisposedException.ThrowIf(released, this);
    internal void Register(object owner, IWeakStorage storage)
    {
        EnsureActive();
        registrations.Add(new(owner, storage)); ++tableCreated;
    }
    internal void Created() => ++created;
    internal void Disposed() => ++disposed;
    public WeakCollectionLifetime Prune()
    {
        int entries = 0;
        for (int index = registrations.Count - 1; index >= 0; --index)
        {
            var item = registrations[index];
            if (item.Owner.Target is null)
            {
                item.Storage.Clear(); item.Owner.Free(); ++tableDisposed;
                registrations.RemoveAt(index);
            }
            else entries += item.Storage.Prune();
        }
        return new(registrations.Count, entries, created, disposed, tableCreated, tableDisposed);
    }
    public void Release()
    {
        if (released) return;
        foreach (var item in registrations)
        { item.Storage.Clear(); item.Owner.Free(); ++tableDisposed; }
        registrations.Clear(); released = true;
    }
    public WeakCollectionLifetime Lifetime => new(registrations.Count,
        registrations.Sum(item => item.Storage.Count), created, disposed, tableCreated, tableDisposed);
}

internal interface IWeakStorage { int Count { get; } int Prune(); void Clear(); }

/// <summary>Storage-only replacement for Jint's original ConditionalWeakTable.
/// Jint retains its intrinsic constructors, prototypes, validation and algorithms.</summary>
public sealed class OwnedWeakTable<TKey,TValue> where TKey : class where TValue : class
{
    private sealed class Storage(WeakCollectionScope scope) : IWeakStorage
    {
        internal sealed class Entry(TKey key, TValue value) { internal DependentHandle Handle = new(key, value); }
        internal readonly Dictionary<int,List<Entry>> Buckets = [];
        private readonly List<int> empty = [];
        public int Count { get; private set; }
        internal void Added() { ++Count; scope.Created(); }
        internal void Free(Entry entry) { entry.Handle.Dispose(); --Count; scope.Disposed(); }
        public int Prune()
        {
            foreach (var item in Buckets)
            {
                var bucket = item.Value;
                for (int index = bucket.Count - 1; index >= 0; --index)
                    if (bucket[index].Handle.Target is null) { Free(bucket[index]); bucket.RemoveAt(index); }
                if (bucket.Count == 0) empty.Add(item.Key);
            }
            foreach (int hash in empty) Buckets.Remove(hash);
            empty.Clear(); return Count;
        }
        public void Clear()
        {
            foreach (var bucket in Buckets.Values) foreach (var entry in bucket) Free(entry);
            Buckets.Clear(); empty.Clear();
        }
    }
    private readonly Storage storage;
    private readonly WeakCollectionScope scope;
    public OwnedWeakTable(WeakCollectionScope scope)
    { ArgumentNullException.ThrowIfNull(scope); this.scope=scope; storage = new(scope); scope.Register(this, storage); }
    public bool TryGetValue(TKey key, out TValue value)
    {
        scope.EnsureActive(); ArgumentNullException.ThrowIfNull(key);
        if (storage.Buckets.TryGetValue(RuntimeHelpers.GetHashCode(key), out var bucket))
            foreach (var entry in bucket)
            {
                var pair = entry.Handle.TargetAndDependent;
                if (ReferenceEquals(pair.Target, key)) { value = (TValue)pair.Dependent!; return true; }
            }
        value = null!; return false;
    }
    public void Add(TKey key, TValue value) => Store(key, value, false);
    public void AddOrUpdate(TKey key, TValue value) => Store(key, value, true);
    // Keep the original callback type so the interpreter's cached delegates and
    // its lazy shape factory remain unchanged. A factory may re-enter this table.
    public TValue GetValue(TKey key, ConditionalWeakTable<TKey,TValue>.CreateValueCallback createValueCallback)
    {
        scope.EnsureActive(); ArgumentNullException.ThrowIfNull(createValueCallback);
        if (TryGetValue(key, out var present)) return present;
        var created = createValueCallback(key);
        if (TryGetValue(key, out present)) return present;
        Add(key, created); return created;
    }
    // The pinned private-element store supplies its concrete dictionary factory;
    // this avoids reflection-based construction on NativeAOT.
    public TValue GetOrCreateValue(TKey key, Func<TValue> factory)
    {
        scope.EnsureActive(); ArgumentNullException.ThrowIfNull(factory);
        if (TryGetValue(key, out var present)) return present;
        var created = factory();
        if (TryGetValue(key, out present)) return present;
        Add(key, created); return created;
    }
    private void Store(TKey key, TValue value, bool update)
    {
        scope.EnsureActive(); ArgumentNullException.ThrowIfNull(key);
        int hash = RuntimeHelpers.GetHashCode(key);
        if (!storage.Buckets.TryGetValue(hash, out var bucket)) storage.Buckets.Add(hash, bucket = []);
        for (int index = bucket.Count - 1; index >= 0; --index)
        {
            var entry = bucket[index]; var pair = entry.Handle.TargetAndDependent;
            if (pair.Target is null) { storage.Free(entry); bucket.RemoveAt(index); continue; }
            if (!ReferenceEquals(pair.Target, key)) continue;
            if (!update) throw new ArgumentException("The weak identity key already exists.", nameof(key));
            entry.Handle.Dependent = value; return;
        }
        bucket.Add(new(key, value)); storage.Added();
    }
    public bool Remove(TKey key)
    {
        scope.EnsureActive(); ArgumentNullException.ThrowIfNull(key);
        int hash = RuntimeHelpers.GetHashCode(key);
        if (storage.Buckets.TryGetValue(hash, out var bucket))
            for (int index = 0; index < bucket.Count; ++index)
            {
                var entry = bucket[index];
                if (!ReferenceEquals(entry.Handle.Target, key)) continue;
                storage.Free(entry); bucket.RemoveAt(index);
                if (bucket.Count == 0) storage.Buckets.Remove(hash); return true;
            }
        return false;
    }
}