using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;

namespace AlicizaX
{

    public static class MemoryPoolRegistry
    {
        internal sealed class MemoryPoolHandle
        {
            public readonly Func<IMemory> Acquire;
            public readonly Action<IMemory> Release;
            public readonly Action Clear;
            public readonly Action<int> Prewarm;
            public readonly Func<MemoryPoolInfo> GetInfo;
            public readonly Action<int> Tick;
            public readonly Action<int> Shrink;

            public MemoryPoolHandle(
                Func<IMemory> acquire,
                Action<IMemory> release,
                Action clear,
                Action<int> prewarm,
                Func<MemoryPoolInfo> getInfo,
                Action<int> tick,
                Action<int> shrink)
            {
                Acquire = acquire;
                Release = release;
                Clear = clear;
                Prewarm = prewarm;
                GetInfo = getInfo;
                Tick = tick;
                Shrink = shrink;
            }
        }

        private static readonly Dictionary<Type, MemoryPoolHandle> s_Handles
            = new Dictionary<Type, MemoryPoolHandle>(64);

        // Tick 回调用数组缓存，避免每帧遍历 Dictionary
        private static Action<int>[] s_TickArray = Array.Empty<Action<int>>();
        private static int s_TickCount;
        private static bool s_TickArrayDirty;

        public static int Count => s_Handles.Count;

        internal static void Register(Type type, MemoryPoolHandle handle)
        {
            s_Handles[type] = handle;
            s_TickArrayDirty = true;
        }

        /// <summary>
        /// 非泛型 Acquire，用于只有 Type 没有泛型参数的场景。
        /// </summary>
        public static IMemory Acquire(Type type)
        {
            if (s_Handles.TryGetValue(type, out var handle))
                return handle.Acquire();

            RuntimeHelpers.RunClassConstructor(
                typeof(MemoryPool<>).MakeGenericType(type).TypeHandle);

            if (s_Handles.TryGetValue(type, out handle))
                return handle.Acquire();

            throw new Exception($"MemoryPool: Type '{type.FullName}' is not a valid IMemory type.");
        }

        public static void Release(IMemory memory)
        {
            if (memory == null)
                throw new ArgumentNullException(nameof(memory));

            Type type = memory.GetType();
            if (s_Handles.TryGetValue(type, out var handle))
            {
                handle.Release(memory);
                return;
            }

            RuntimeHelpers.RunClassConstructor(
                typeof(MemoryPool<>).MakeGenericType(type).TypeHandle);

            if (s_Handles.TryGetValue(type, out handle))
                handle.Release(memory);
        }

        public static MemoryPoolInfo[] GetAllInfos()
        {
            var infos = new MemoryPoolInfo[s_Handles.Count];
            int i = 0;
            foreach (var kv in s_Handles)
                infos[i++] = kv.Value.GetInfo();
            return infos;
        }

        public static void ClearAll()
        {
            foreach (var kv in s_Handles)
                kv.Value.Clear();
        }

        public static void Prewarm(Type type, int count)
        {
            if (s_Handles.TryGetValue(type, out var handle))
            {
                handle.Prewarm(count);
                return;
            }

            RuntimeHelpers.RunClassConstructor(
                typeof(MemoryPool<>).MakeGenericType(type).TypeHandle);

            if (s_Handles.TryGetValue(type, out handle))
                handle.Prewarm(count);
        }

        public static void ClearType(Type type)
        {
            if (s_Handles.TryGetValue(type, out var handle))
                handle.Clear();
        }

        public static void RemoveFromType(Type type, int count)
        {
            if (s_Handles.TryGetValue(type, out var handle))
            {
                int unused = handle.GetInfo().UnusedCount;
                handle.Shrink(unused - count);
            }
        }


        public static void TickAll(int frameCount)
        {
            if (s_TickArrayDirty)
                RebuildTickArray();

            for (int i = 0; i < s_TickCount; i++)
                s_TickArray[i](frameCount);
        }

        private static void RebuildTickArray()
        {
            s_TickCount = s_Handles.Count;
            if (s_TickArray.Length < s_TickCount)
                s_TickArray = new Action<int>[s_TickCount];

            int i = 0;
            foreach (var kv in s_Handles)
                s_TickArray[i++] = kv.Value.Tick;

            s_TickArrayDirty = false;
        }
    }
}