namespace UnmanagedMMU
{
    using System;
    using System.Collections.Generic;
    using System.Runtime.CompilerServices;
    using UnmanagedMMU.Allocators;
    using UnmanagedMMU.Handles;
    using UnmanagedMMU.Handles.Internal;

    /// <summary>
    /// Provides an unmanaged heap for long-lived allocations with reuse.
    /// </summary>
    /// <remarks>
    /// <para>
    /// <see cref="WorkspaceHeap"/> minimizes calls to the underlying allocator by retaining
    /// freed blocks in size-segregated free lists and reusing them when possible.
    /// </para>
    /// <para>
    /// Allocation strategy:
    /// <list type="bullet">
    /// <item><description>
    /// <b>Small allocations</b> (≤ 1 KB) use fixed size buckets
    /// </description></item>
    /// <item><description>
    /// <b>Medium allocations</b> (≤ 256 KB) use best-fit reuse
    /// </description></item>
    /// <item><description>
    /// <b>Large allocations</b> (&gt; 256 KB) use tolerance-based reuse
    /// </description></item>
    /// </list>
    /// </para>
    /// </remarks>
    public unsafe sealed class WorkspaceHeap : IDisposable, IUnmanagedMemoryOwner
    {

        /// <summary>
        /// Struct defining metadata to be stored with each allocation
        /// </summary>
        private struct AllocationMetadata
        {
            public IntPtr Block;
            public nuint Alignment;
        }

        /// <summary>
        /// The maximum size, in bytes, for "small" allocations. Uses fixed-size buckets
        /// </summary>
        private const nuint _smallThreshold = 1024;

        /// <summary>
        /// The maximum size, in bytes, for allocations considered "medium". Uses best-fit reuse
        /// </summary>
        private const nuint _mediumThreshold = 256 * 1024; // 256 KB

        /// <summary>
        /// The maximum absolute number of bytes that may be wasted when reusing a large block for a "large" allocation.
        /// </summary>
        private const nuint _largeMaxWasteBytes = 256 * 1024; // 256 KB

        /// <summary>
        /// The maximum allowed size ratio when reusing a large allocation block.
        /// </summary>
        /// <remarks>
        /// For example, a value of <c>1.25</c> allows a block up to 25% larger than
        /// the requested size to be reused.
        /// </remarks>
        private const double _largeWasteRatioLimit = 1.25;

        /// <summary>
        /// Predefined bucket sizes used for small allocation reuse.
        /// </summary>
        private static readonly nuint[] _sizeClasses =
        [
             32,  64,  96, 128, 160, 192, 224, 256,
            288, 320, 352, 384, 416, 448, 480, 512,
            544, 576, 608, 640, 672, 704, 736, 768,
            800, 832, 864, 896, 928, 960, 992, 1024
        ];

        /// <summary>
        /// Allocator interface used for all underlying unmanaged memory operations.
        /// </summary>
        private readonly IUnmanagedAllocator _allocator;

        /// <summary>
        /// Internal lock, ensures thread safety while maintaining a simple interface
        /// </summary>
        private readonly Lock _lock = new();

        /// <summary>
        /// Free lists for small allocations, keyed by the bucket size, in bytes.
        /// </summary>
        /// <remarks>
        /// For a given bucket, the corresponding stack contains previously allocated blocks that are available to be used
        /// </remarks>
        private readonly Dictionary<nuint, List<AllocationMetadata>> _smallFree = [];

        /// <summary>
        /// Free lists for medium allocations keyed by the exactly allocated size, in bytes, and sorted for best-fit.
        /// </summary>
        private readonly SortedDictionary<nuint, List<AllocationMetadata>> _mediumFree = [];

        /// <summary>
        /// Free lists for large allocations keyed by the exactly allocated size, in bytes, sorted for tolerance-based reuse.
        /// </summary>
        private readonly SortedDictionary<nuint, List<AllocationMetadata>> _largeFree = [];

        /// <summary>
        /// Tracks the total bytes of memory reserved from the provided <see cref="IUnmanagedAllocator"/>.
        /// </summary>
        private nuint _totalReserved;

        /// <summary>
        /// Total memory currently in use by active allocations, in bytes.
        /// </summary>
        private nuint _totalInUse;

        /// <summary>
        /// Counts actual underlying OS allocations.
        /// </summary>
        private nuint _totalAllocations;

        /// <summary>
        /// Indicates whether this <see cref="WorkspaceHeap"/> has been disposed.
        /// </summary>
        private volatile bool _disposed;

        /// <summary>
        /// Creates a new <see cref="WorkspaceHeap"/>.
        /// </summary>
        public WorkspaceHeap()
            : this(new DefaultUnmanagedAllocator())
        {
        }

        /// <summary>
        /// Creates a new <see cref="WorkspaceHeap"/> using the specified  <see cref="IUnmanagedAllocator"/>.
        /// </summary>
        /// <param name="allocator">Allocator implementing <see cref="IUnmanagedAllocator"/>.</param>
        internal WorkspaceHeap(IUnmanagedAllocator allocator)
        {
            _allocator = allocator;

            // Initialize small-size buckets
            foreach (nuint size in _sizeClasses)
            {
                _smallFree[size] = [];                    
            }
        }

        /// <summary>
        /// Gets the total number of bytes currently allocated from the underlying allocator.
        /// </summary>
        /// <remarks>
        /// This includes both active allocations and freed blocks retained for reuse.
        /// </remarks>
        public nuint TotalReservedBytes
        {
            get
            {
                lock (_lock)
                    return _totalReserved;
            }
        }

        /// <summary>
        /// Gets the total number of bytes currently in use by active allocations.
        /// </summary>
        /// <remarks>
        /// This value decreases when memory is freed and increases when new allocations occur.
        /// </remarks>
        public nuint TotalUsedBytes
        {
            get
            {
                lock (_lock)
                {
                    return _totalInUse;                    
                }
            }
        }

        /// <summary>
        /// Gets the total number of allocation operations performed by this heap.
        /// </summary>
        /// <remarks>
        /// This counts new underlying OS allocations, not reuse from free lists.
        /// Useful for performance diagnostics and testing reuse behavior.
        /// </remarks>
        public nuint TotalAllocationCount
        {
            get
            {
                lock (_lock)
                {
                    return _totalAllocations;                    
                }
            }
        }

        /// <summary>
        /// Indicates whether the <see cref="WorkspaceHeap"/> has been disposed.
        /// </summary>
        public bool IsDisposed
        {
            get { return _disposed; }
        }

        /// <summary>
        /// Determines the small size bucket for a requested allocation.
        /// </summary>
        /// <param name="size">The allocation size to get the bucket size for</param>
        /// <returns>The small size bucket for the requested allocation</returns>
        private static nuint GetSizeClass(nuint size)
        {
            foreach (nuint s in _sizeClasses)
            {
                if (size <= s)
                {
                    return s;
                }
            }
            return size;
        }

        /// <summary>
        /// Allocates a new block from the underlying allocator including a header.
        /// </summary>
        /// <param name="payloadSize">Requested payload size in bytes.</param>
        /// <param name="alignment">FOO</param>
        /// <returns>
        /// Pointer to the allocated block (header included).
        /// </returns>
        private IntPtr AllocateNewAligned(nuint payloadSize, nuint alignment)
        {
            void* raw = _allocator.AllocAligned(payloadSize, alignment);
            _totalReserved += payloadSize;
            _totalAllocations++;
            return (IntPtr)raw;
        }

        /// <summary>
        /// Allocates unmanaged memory from the workspace heap.
        /// </summary>
        /// <param name="count">Number of elements <typeparamref name="T"/> to allocate.</param>
        /// <param name="zero">If true, memory is zero-initialized.</param>
        /// <returns> An <see cref="IMemoryHandle{T}"/> to the allocated memory.</returns>
        /// <exception cref="ObjectDisposedException">Thrown if heap is disposed.</exception>
        /// <exception cref="ArgumentOutOfRangeException">Thrown if size is zero.</exception>
        /// <remarks>The memory returned by this method is always 16-byte aligned, for other alignemnts use </remarks>
        public IMemoryHandle<T> Allocate<T>(int count, bool zero = false) where T : unmanaged
        {
            return AllocateAligned<T>(count, SegmentAlignment.Aligned16, zero);
        }

        /// <summary>
        /// 
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="count"></param>
        /// <param name="alignment"></param>
        /// <param name="zero"></param>
        /// <returns></returns>
        /// <exception cref="OverflowException"></exception>
        public IMemoryHandle<T> AllocateAligned<T>(int count, SegmentAlignment alignment, bool zero = false) where T: unmanaged
        {
            ArgumentOutOfRangeException.ThrowIfNegative(count);
            ArgumentOutOfRangeException.ThrowIfZero(count);
            if ((nuint)count > nuint.MaxValue / (nuint)(sizeof(T)))
            {
                throw new OverflowException($"Requested allocation of {count} elements of type {typeof(T)} exceeds allowable maximum memory size.");
            }

            nuint size = (nuint)count * (nuint)sizeof(T);
            nuint requestedAlignment = (nuint)alignment;
            nuint effectiveAlignment = requestedAlignment < (nuint)sizeof(T) ? (nuint)sizeof(T) : requestedAlignment;

            lock (_lock)
            {
                ThrowIfDisposed();
                void* ptr;

                if (size <= _smallThreshold)
                    ptr = AllocateSmallAligned(size, effectiveAlignment, zero);
                else if (size <= _mediumThreshold)
                    ptr = AllocateMediumAligned(size, effectiveAlignment, zero);
                else
                    ptr = AllocateLargeAligned(size, effectiveAlignment, zero);

                return new PersistentMemoryHandle<T>((T*)ptr, size, effectiveAlignment, this);
            }
        }

        /// <summary>
        /// TODO: Fill in
        /// </summary>
        /// <param name="alignment"></param>
        /// <param name="list"></param>
        /// <returns></returns>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static int FindIndexAlignmentMatch(nuint alignment, List<AllocationMetadata> list)
        {
            // Attempt to find an index match the requested size
            int mIdx = -1;
            for (int i = 0; i < list.Count; i++)
            {
                if (list[i].Alignment >= alignment)
                {
                    mIdx = i;
                    break;
                }
            }
            return mIdx;
        }

        /// <summary>
        /// Allocates a small-size block using bucketed free lists.
        /// </summary>
        private void* AllocateSmallAligned(nuint size, nuint alignment, bool zero)
        {
            nuint bucket = GetSizeClass(size);
            List<AllocationMetadata> free = _smallFree[bucket];

            // Attempt to find an index match the requested size
            int mIdx = FindIndexAlignmentMatch(alignment, free);

            if (mIdx >= 0)
            {
                AllocationMetadata meta = free[mIdx];
                free.RemoveAt(mIdx);

                _totalInUse += bucket;
                void* user = (void*)meta.Block;
                if (zero)
                {
                    Unsafe.InitBlockUnaligned(user, 0, (uint)bucket);
                }
                return user;
            }

            // allocate a new small block
            IntPtr block = AllocateNewAligned(bucket, alignment);
            _totalInUse += bucket;
            void* nUser = (void*)block;
            if (zero)
            {
                Unsafe.InitBlockUnaligned(nUser, 0, (uint)bucket);
            }
            return nUser;
        }

        /// <summary>Allocates a medium-size block using best-fit reuse.</summary>
        private void* AllocateMediumAligned(nuint size, nuint alignment, bool zero)
        {
            foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in _mediumFree)
            {
                if (kv.Key >= size && kv.Value.Count > 0)
                {
                    List<AllocationMetadata> free = kv.Value;

                    // Attempt to find an index match the requested size
                    int mIdx = FindIndexAlignmentMatch(alignment, free);

                    if (mIdx >= 0)
                    {
                        AllocationMetadata meta = free[mIdx];
                        free.RemoveAt(mIdx);

                        _totalInUse += kv.Key;
                        void* user = (void*)meta.Block;
                        if (zero)
                        {
                            Unsafe.InitBlockUnaligned(user, 0, (uint)kv.Key);
                        }
                        return user;
                    }
                }
            }
            // no match
            IntPtr newBlock = AllocateNewAligned(size, alignment);
            _totalInUse += size;

            void* nUser = (void*)newBlock;
            if (zero)
            {
                Unsafe.InitBlockUnaligned(nUser, 0, (uint)size);
            }               

            return nUser;
        }

        /// <summary>
        /// Allocates a large block using tolerance-based reuse (smallest ≥ requested within waste bounds).
        /// </summary>
        private void* AllocateLargeAligned(nuint size, nuint alignment, bool zero)
        {
            foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in _largeFree)
            {
                nuint blockSize = kv.Key;
                if (blockSize < size || kv.Value.Count == 0)
                {
                    continue;
                }

                // check waste tolerance before attempting to find an alignment match
                nuint waste = blockSize - size;
                bool acceptable = waste <= _largeMaxWasteBytes || ((double)blockSize / size) <= _largeWasteRatioLimit;

                if (!acceptable)
                {
                    continue;
                }

                List<AllocationMetadata> free = kv.Value;

                // Attempt to find an index match the requested size
                int mIdx = FindIndexAlignmentMatch(alignment, free);

                if (mIdx >= 0)
                {
                    AllocationMetadata meta = free[mIdx];
                    free.RemoveAt(mIdx);

                    _totalInUse += blockSize;
                    void* user = (void*)meta.Block;
                    if (zero)
                    {
                        Unsafe.InitBlockUnaligned(user, 0, (uint)blockSize);
                    }

                    return user;
                }
            }

            // no match

            IntPtr block = AllocateNewAligned(size, alignment);
            _totalInUse += size;
            void* nUser = (void*)block;
            if (zero)
            {
                Unsafe.InitBlockUnaligned(nUser, 0, (uint)size);
            }
            return nUser;
        }

        /// <summary>
        /// Frees a previously allocated block, returning it to the appropriate free list.
        /// </summary>
        void IUnmanagedMemoryOwner.Free(IOwnedHandle handle)
        {
            ThrowIfDisposed();
            if (handle.Pointer == null)
            {
                return;
            }

            if (handle.GetOwner() != this)
            {
                throw new InvalidOperationException(
                    "Attempted to free a handle from a different allocator pool.");
            }

            lock (_lock)
            {
                nuint size = handle.ByteCount;
                nuint alignment = handle.Alignment;
                _totalInUse -= size;

                if (size <= _smallThreshold)
                {
                    _smallFree[size].Add(new AllocationMetadata
                    {
                        Block = (IntPtr)handle.Pointer,
                        Alignment = alignment
                    });
                }
                else if (size <= _mediumThreshold)
                {
                    if (!_mediumFree.TryGetValue(size, out List<AllocationMetadata>? free))
                    {
                        _mediumFree[size] = free = [];
                    }
                    free.Add(new AllocationMetadata
                    {
                        Block = (IntPtr)handle.Pointer,
                        Alignment = alignment
                    });
                }
                else
                {
                    if (!_largeFree.TryGetValue(size, out List<AllocationMetadata>? free))
                    {
                        _largeFree[size] = free = [];
                    }
                    free.Add(new AllocationMetadata
                    {
                        Block = (IntPtr)handle.Pointer,
                        Alignment = alignment
                    });
                }
            }
        }

        /// <summary>
        /// Releases all unused blocks back to the underlying allocator.
        /// </summary>
        public void Prune()
        {
            ThrowIfDisposed();
            lock (_lock)
            {
                PruneDictionary(_smallFree);
                PruneDictionary(_mediumFree);
                PruneDictionary(_largeFree);
            }
        }

        /// <summary>Helper to free all blocks in a dictionary of free stacks.</summary>
        private void PruneDictionary(IDictionary<nuint, List<AllocationMetadata>> dict)
        {
            foreach (KeyValuePair<nuint, List<AllocationMetadata>> kv in dict)
            {
                List<AllocationMetadata> list = kv.Value;
                while (list.Count > 0)
                {
                    var item = list[^1];
                    list.RemoveAt(list.Count - 1);
                    _allocator.FreeAligned((void*)item.Block, item.Alignment);
                    _totalReserved -= kv.Key;
                }
            }
        }

        /// <summary>
        /// Releases all memory and marks the heap as disposed.
        /// </summary>
        public void Dispose()
        {
            if (_disposed)
            {
                return;
            }

            lock (_lock)
            {
                if (_disposed)
                {
                    return;
                }
                if (_totalInUse > 0)
                {
                    throw new InvalidOperationException(
                        "Cannot dispose WorkspaceHeap while active allocations exist. " +
                        "Dispose all handles returned from this heap before disposing the heap.");
                }
                Prune();
                // Reset stats
                _totalInUse = 0;
                _totalAllocations = 0;
                _disposed = true;
            }
        }

        /// <summary>
        /// Throws an <see cref="ObjectDisposedException"/> if the <see cref="WorkspaceHeap"/> has already been disposed.
        /// </summary>
        /// <exception cref="ObjectDisposedException">
        /// Thrown when this instance is no longer valid for use.
        /// </exception>
        private void ThrowIfDisposed()
        {
            ObjectDisposedException.ThrowIf(_disposed, this);
        }
    }
}