namespace UnmanagedMMU
{
    using System;
    using System.Collections.Generic;
    using System.Runtime.InteropServices;
    using System.Threading;
    using UnmanagedMMU.Allocators;

    /// <summary>
    /// Implementation of segmented Bump-Allocator.
    /// 
    /// This implementation manages fixed-sized unmanaged memory segments.
    /// This ensures allocations are fast and contiguous within a segment.
    /// Once a segment is full, a new one is automatically allocated
    /// </summary>
    public unsafe sealed class SegmentedPool : IDisposable
    {
        /// <summary>
        /// The default size for a <see cref="Segment"/>
        /// </summary>
        private const nuint _defaultSegmentSize = 4194304; // 4 MiB

        /// <summary>
        /// The size that a new <see cref="Segment"/> should be
        /// </summary>
        private nuint _currentSegmentSize;

        /// <summary>
        /// Queue of free segments
        /// </summary>
        private readonly Stack<IntPtr> _freeSegments = new();

        /// <summary>
        /// List of active memory segments
        /// </summary>
        private readonly List<IntPtr> _activeSegments = [];

        /// <summary>
        /// Tracks the total amount of allocated bytes
        /// </summary>
        private nuint _totalAllocated = 0;

        /// <summary>
        /// Tracks the total amount of allocated bytes current in use
        /// </summary>
        private nuint _totalUsed = 0;

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

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

        /// <summary>
        /// Pointer to the currently in use <see cref="Segment"/>
        /// </summary>
        private Segment* _current;

        private readonly IUnmanagedAllocator _allocator;

        /// <summary>
        /// Represents a memory segment in the <see cref="SegmentedPool"/>.
        /// </summary>
        /// <remarks>
        /// <list type="bullet">
        ///   <item><description><see cref="Ptr"/> Pointer to the start of the unmanaged memory block</description></item>
        ///   <item><description><see cref="Offset"/> indicates the current allocation position within the segment. Each new allocation advances this offset by the number of bytes allocated.</description></item>
        ///   <item><description><see cref="Size"/> is the total size, in bytes, of the segment.</description></item>
        ///   <item><description>Segments are managed internally by the <see cref="SegmentedPool"/> and should not be modified directly outside of the pool.</description></item>
        ///   <item><description>A <see cref="Segment"/> is a contiguous block of unmanaged memory from which allocations are served sequentially via bump allocation.</description></item>
        /// </list>
        /// </remarks>
        private struct Segment
        {
            /// <summary>
            /// Pointer to the start of the unmanaged memory block.
            /// </summary>
            public byte* Ptr;

            /// <summary>
            /// The current offset into the <see cref="Segment"/> where the next allocation will occur.
            /// </summary>
            public nuint Offset;

            /// <summary>
            /// Total size of the  <see cref="Segment"/> in bytes.
            /// </summary>
            public nuint Size;
        }


        /// <summary>
        /// Initializes a new <see cref="SegmentedPool"/> with the specified default <paramref name="segmentSize"/> and <paramref name="initialSegments"/> count
        /// </summary>
        /// <param name="segmentSize">Size of each segment in bytes (default 4 MiB)</param>
        /// <param name="initialSegments">Number of segments to pre-allocate to the pool</param>
        /// <exception cref="ArgumentException">
        /// Thrown if <paramref name="segmentSize"/> is zero, or if <paramref name="initialSegments"/> is less than 1.
        /// </exception>
        public SegmentedPool(nuint segmentSize = _defaultSegmentSize, int initialSegments = 4)
            : this(segmentSize, initialSegments, new DefaultUnmanagedAllocator())
        {
        }

        /// <summary>
        /// Initializes a new <see cref="SegmentedPool"/> with the specified default <paramref name="segmentSize"/> and <paramref name="initialSegments"/> count
        /// </summary>
        /// <param name="segmentSize">Size of each segment in bytes (default 4 MiB)</param>
        /// <param name="initialSegments">Number of segments to pre-allocate to the pool</param>
        /// <param name="allocator">IUnmanagedAllocator instance that implements the allocator</param>
        /// <exception cref="ArgumentException">
        /// Thrown if <paramref name="segmentSize"/> is zero, or if <paramref name="initialSegments"/> is less than 1.
        /// </exception>
        internal SegmentedPool(nuint segmentSize, int initialSegments, IUnmanagedAllocator allocator)
        {
            if (segmentSize == 0)
            {
                throw new ArgumentException("Segment size must be greater than zero.", nameof(segmentSize));
            }
            if (initialSegments < 1)
            {
                throw new ArgumentException("Initial segments count must be at least 1.", nameof(initialSegments));
            }
            _allocator = allocator;
            _currentSegmentSize = segmentSize;

            // Pre-allocate segments
            for (int i = 0; i < initialSegments; i++)
            {
                _freeSegments.Push((IntPtr)AllocateNewSegment(_currentSegmentSize));
            }
                

            _current = (Segment*)_freeSegments.Pop();
            _activeSegments.Add((IntPtr)_current);
        }

        /// <summary>
        /// Gets the number of free segments  available for use in the pool.
        /// </summary>
        /// <returns>The number of free segments.</returns>
        public int FreeSegmentCount
        {
            get { return _freeSegments.Count; }
        }

        /// <summary>
        /// Gets the number of segments currently in use in the pool.
        /// </summary>
        /// <returns>The number of currently active Segments.</returns>
        public int ActiveSegmentCount
        {
            get { return _activeSegments.Count; }
        }

        /// <summary>
        /// Returns the total number of bytes that have currently been allocated.
        /// </summary>
        /// <returns>The total number of bytes that have been allocated.</returns>
        public nuint TotalAllocatedBytes
        {
            get { return _totalAllocated; }
        }

        /// <summary>
        /// Gets the total number of bytes currently in use across all active segments.
        /// </summary>
        /// <returns>The total number of bytes that are in use.</returns>
        public nuint TotalUsedBytes
        {
            get { return _totalUsed; }
        }

        /// <summary>
        /// Gets the size, in bytes, used when allocating new <see cref="Segment"/> instances.
        /// </summary>
        /// <remarks>
        /// This reflects the most recently configured size and affects only future <see cref="Segment"/> allocations.
        /// </remarks>
        public nuint CurrentSegmentSize
        {
            get { return _currentSegmentSize; }
        }

        /// <summary>
        /// Gets a value indicating whether this <see cref="SegmentedPool"/> instance has been disposed.
        /// </summary>
        /// <returns>True, if the current <see cref="SegmentedPool"/> instance has been disposed of, false otherwise.</returns>
        /// <remarks>
        /// Once disposed, further calls to allocation or reset methods will throw <see cref="ObjectDisposedException"/>.
        /// </remarks>
        public bool IsDisposed
        {
            get { return _disposed; }
        }

        /// <summary>
        /// Sets the current <see cref="Segment"/> size used for subsequent allocations.
        /// </summary>
        /// <param name="newSize">
        /// The new segment size, in bytes, that will be used when allocating future <see cref="Segment"/> instances.
        /// Must be greater than zero.
        /// </param>
        /// <remarks>
        /// This does not affect any <see cref="Segment"/> that have already been allocated or are currently active.
        /// Only new <see cref="Segment"/> created after calling this method will use the updated size.
        /// </remarks>
        /// <exception cref="ArgumentOutOfRangeException">
        /// Thrown if <paramref name="newSize"/> is zero.
        /// </exception>
        public void SetSegmentSize(nuint newSize)
        {
            ThrowIfDisposed();
            if (newSize == 0)
            {
                throw new ArgumentOutOfRangeException(nameof(newSize), "Segment size must be greater than zero.");
            }
            _currentSegmentSize = newSize;
        }

        /// <summary>
        /// Resets the current <see cref="Segment"/> size back to the default 4 MiB
        /// </summary>
        /// <remarks>
        /// Future <see cref="Segment"/> allocations will revert to using this default size 
        /// Active and free <see cref="Segment"/> are not modified.
        /// </remarks>
        public void ResetSegmentSize()
        {
            ThrowIfDisposed();
            _currentSegmentSize = _defaultSegmentSize;
        }

        /// <summary>
        /// Allocates a span of unmanaged memory of size <paramref name="count"/> for elements of type <typeparamref name="T"/>.
        /// </summary>
        /// <typeparam name="T">The unmanaged value type to store in the allocated memory. Must be a struct or primitive type.</typeparam>
        /// <param name="count">The number of elements of type <typeparamref name="T"/> to allocate.</param>
        /// <returns>A <see cref="Span{T}"/> representing the allocated memory. The span is valid until the <see cref="SegmentedPool"/> is reset or disposed.</returns>
        /// <remarks>
        /// <list type="bullet">
        ///   <item><description>This allocation is performed in unmanaged memory and bypasses the .NET garbage collector.</description></item>
        ///   <item><description>Accessing the memory after <see cref="Reset"/> or <see cref="Dispose"/> has been called is undefined behavior and may lead to crashes.</description></item>
        /// </list>
        /// </remarks>
        /// <exception cref="ArgumentOutOfRangeException">
        /// Thrown if <paramref name="count"/> is less than or equal to zero.
        /// </exception>
        /// <exception cref="OverflowException">
        /// Thrown if the total allocation size (count * sizeof(T)) exceeds the maximum allowable size.
        /// </exception>
        public Span<T> Allocate<T>(int count) where T : unmanaged
        {
            ThrowIfDisposed();
            if (count <= 0)
            {
                throw new ArgumentOutOfRangeException(nameof(count), "Allocation count must be greater than zero.");
            }
            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 bytes = (nuint)(count * sizeof(T));

            lock (_lock)
            {
                // Enough space in current segment?
                if (_current->Offset + bytes > _current->Size)
                    SwitchSegment(bytes);

                T* ptr = (T*)(_current->Ptr + _current->Offset);
                _current->Offset += bytes;
                _totalUsed += bytes;
                return new Span<T>(ptr, count);
            }
        }

        /// <summary>
        /// Switches to a new <see cref="Segment"/> when the current <see cref="Segment"/> is full.
        /// </summary>
        /// <param name="requiredBytes">The number of bytes required for the upcoming allocation. If the current <see cref="Segment"/> does not have enough free space, a new <see cref="Segment"/> will be used.</param>
        private void SwitchSegment(nuint requiredBytes)
        {
            Segment* segment;

            // Allocate fresh Segment if needed
            if (_freeSegments.Count == 0 || requiredBytes > _currentSegmentSize)
            {
                segment = AllocateNewSegment(requiredBytes > _currentSegmentSize ? requiredBytes : _currentSegmentSize);
            }
            else
            {
                segment = (Segment*)_freeSegments.Pop();
                segment->Offset = 0;
            }

            _activeSegments.Add((IntPtr)segment);
            _current = segment;
        }

        /// <summary>
        /// Allocates a new <see cref="Segment"/>
        /// </summary>
        /// <param name="size">
        /// Optional size, in bytes, for the new <see cref="Segment"/>. 
        /// If <c>null</c>, the default <see cref="Segment"/> size (<see cref="_defaultSegmentSize"/>) is used.
        /// </param>
        /// <returns>A pointer to the newly allocated <see cref="Segment"/></returns>
        private Segment* AllocateNewSegment(nuint size)
        {
            byte* ptr = (byte*)_allocator.Alloc(size);
            Segment* segment = (Segment*)_allocator.Alloc((nuint)sizeof(Segment));
            segment->Ptr = ptr;
            segment->Offset = 0;
            segment->Size = size;

            _totalAllocated += size;
            return segment;
        }

        /// <summary>
        /// Frees unused <see cref="Segment"/> in the free pool, reducing unmanaged memory usage.
        /// </summary>
        /// <param name="minFreeSegments"> Minimum number of free segments to retain for future allocations. Defaults to <c>16</c>. </param>
        /// <remarks>
        /// Segments beyond the retained count will have their unmanaged memory released back to the system.
        /// </remarks>
        /// <exception cref="ArgumentOutOfRangeException">
        /// Thrown if <paramref name="minFreeSegments"/> is negative.
        /// </exception>
        public void Trim(int minFreeSegments = 16)
        {
            ThrowIfDisposed();
            ArgumentOutOfRangeException.ThrowIfNegative(minFreeSegments, nameof(minFreeSegments));
            lock (_lock)
            {
                while (_freeSegments.Count > minFreeSegments)
                {
                    var ip = _freeSegments.Pop();
                    Segment* segment = (Segment*)ip;

                    // Free the unmanaged memory
                    _allocator.Free(segment->Ptr);
                    _totalAllocated -= segment->Size;
                    _allocator.Free(segment);
                }                
            }
        }

        /// <summary>
        /// Resets the allocator, returning all active segments to the free pool.
        /// </summary>
        /// <param name="trim">If true, trims the <see cref="Segment"/> in the free pool. <see cref="Trim"/></param>
        public void Reset(bool trim = false)
        {
            ThrowIfDisposed();
            lock (_lock)
            {
                foreach (var ip in _activeSegments)
                {
                    Segment* segment = (Segment*)ip;
                    segment->Offset = 0;
                    _freeSegments.Push(ip);
                }

                _activeSegments.Clear();
                _totalUsed = 0;

                if (_freeSegments.Count > 0)
                {
                    _current = (Segment*)_freeSegments.Pop();
                    _activeSegments.Add((IntPtr)_current);
                }
                else
                {
                    // This should not be hit in normal circumstances as we always have _current
                    _current = AllocateNewSegment(_currentSegmentSize);
                    _activeSegments.Add((IntPtr)_current);
                }
                // Optionally trim excess free segments after reset
                if (trim)
                {
                    Trim();
                }
            }            
        }

        /// <summary>
        /// Releases all unmanaged memory allocated by the <see cref="SegmentedPool"/> and clears internal state.
        /// After calling this method, the pool can no longer be used for allocations.
        /// </summary>
        /// <remarks>
        /// After disposal, all allocations become invalid and any further operations will throw <see cref="ObjectDisposedException"/>.
        /// This method is thread-safe and may be called multiple times safely.
        /// </remarks>
        public void Dispose()
        {
            lock (_lock)
            {
                if (_disposed)
                {
                    return;
                }

                // Free active pages
                foreach (var ip in _activeSegments)
                {
                    Segment* segment = (Segment*)ip;
                    _allocator.Free(segment->Ptr);
                    _allocator.Free(segment);
                }

                // Free free pages
                foreach (var ip in _freeSegments)
                {
                    Segment* segment = (Segment*)ip;
                    _allocator.Free(segment->Ptr);
                    _allocator.Free(segment);
                }

                _activeSegments.Clear();
                _freeSegments.Clear();
                _current = null;
                _totalAllocated = 0;
                _totalUsed = 0;
                _disposed = true;
            }
        }

        /// <summary>
        /// Throws an <see cref="ObjectDisposedException"/> if the <see cref="SegmentedPool"/> 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);
        }
    }
}