using System;
using System.Text;
using System.Collections.Concurrent;
using System.Threading;

namespace EonaCat.LogStack.Pooling;

// This file is part of the EonaCat project(s) which is released under the Apache License.
// See the LICENSE file or go to https://EonaCat.com/License for full license details.

/// <summary>
/// High-performance StringBuilder pool with automatic capacity management.
/// Reuses StringBuilders to reduce GC pressure in high-throughput scenarios.
/// </summary>
public sealed class StringBuilderPool : IDisposable
{
    private readonly ConcurrentBag<StringBuilder> _smallPool;   // <= 1KB
    private readonly ConcurrentBag<StringBuilder> _mediumPool;  // 1KB - 8KB
    private readonly ConcurrentBag<StringBuilder> _largePool;   // > 8KB

    private const int SmallCapacity = 1024;
    private const int MediumCapacity = 8192;
    private const int SmallPoolLimit = 20;
    private const int MediumPoolLimit = 10;
    private const int LargePoolLimit = 5;

    private int _smallCount;
    private int _mediumCount;
    private int _largeCount;
    private bool _isDisposed;

    public StringBuilderPool()
    {
        _smallPool = new ConcurrentBag<StringBuilder>();
        _mediumPool = new ConcurrentBag<StringBuilder>();
        _largePool = new ConcurrentBag<StringBuilder>();
    }

    /// <summary>
    /// Rents a StringBuilder from the appropriate pool based on requested capacity
    /// </summary>
    public StringBuilder Rent(int capacity = 1024)
    {
        if (_isDisposed)
        {
            throw new ObjectDisposedException(nameof(StringBuilderPool));
        }

        ConcurrentBag<StringBuilder>? pool;
        int currentCount;

        if (capacity <= SmallCapacity)
        {
            pool = _smallPool;
            currentCount = _smallCount;
        }
        else if (capacity <= MediumCapacity)
        {
            pool = _mediumPool;
            currentCount = _mediumCount;
        }
        else
        {
            pool = _largePool;
            currentCount = _largeCount;
        }

        if (currentCount > 0 && pool.TryTake(out var sb))
        {
            if (capacity <= SmallCapacity)
            {
                Interlocked.Decrement(ref _smallCount);
            }
            else if (capacity <= MediumCapacity)
            {
                Interlocked.Decrement(ref _mediumCount);
            }
            else
            {
                Interlocked.Decrement(ref _largeCount);
            }

            sb.Clear();
            return sb;
        }

        return new StringBuilder(capacity);
    }

    public void Return(StringBuilder sb)
    {
        if (_isDisposed || sb == null)
        {
            return;
        }

        int capacity = sb.Capacity;
        ConcurrentBag<StringBuilder>? pool;
        int maxLimit;

        if (capacity <= SmallCapacity)
        {
            pool = _smallPool;
            maxLimit = SmallPoolLimit;
            if (Interlocked.Increment(ref _smallCount) <= maxLimit)
            {
                sb.Clear();
                pool.Add(sb);
            }
            else
            {
                Interlocked.Decrement(ref _smallCount);
            }
        }
        else if (capacity <= MediumCapacity)
        {
            pool = _mediumPool;
            maxLimit = MediumPoolLimit;
            if (Interlocked.Increment(ref _mediumCount) <= maxLimit)
            {
                sb.Clear();
                pool.Add(sb);
            }
            else
            {
                Interlocked.Decrement(ref _mediumCount);
            }
        }
        else
        {
            pool = _largePool;
            maxLimit = LargePoolLimit;
            if (Interlocked.Increment(ref _largeCount) <= maxLimit)
            {
                sb.Clear();
                pool.Add(sb);
            }
            else
            {
                Interlocked.Decrement(ref _largeCount);
            }
        }
    }

    /// <summary>
    /// Gets pooled item counts for diagnostics
    /// </summary>
    public (int small, int medium, int large) GetPoolCounts()
        => (_smallCount, _mediumCount, _largeCount);

    /// <summary>
    /// Clears all pools
    /// </summary>
    public void Clear()
    {
        while (_smallPool.TryTake(out _))
        {
            ;
        }

        while (_mediumPool.TryTake(out _))
        {
            ;
        }

        while (_largePool.TryTake(out _))
        {
            ;
        }

        _smallCount = 0;
        _mediumCount = 0;
        _largeCount = 0;
    }

    public void Dispose()
    {
        if (_isDisposed)
        {
            return;
        }

        _isDisposed = true;
        Clear();
    }
}

/// <summary>
/// Rented StringBuilder that automatically returns itself to the pool when disposed
/// </summary>
public sealed class PooledStringBuilder : IDisposable
{
    private readonly StringBuilderPool _pool;
    private readonly StringBuilder _sb;

    public PooledStringBuilder(StringBuilderPool pool, int capacity = 1024)
    {
        _pool = pool ?? throw new ArgumentNullException(nameof(pool));
        _sb = pool.Rent(capacity);
    }

    public StringBuilder Builder => _sb;

    public override string ToString() => _sb.ToString();

    public void Dispose()
    {
        _pool.Return(_sb);
    }
}
