using System.Runtime.CompilerServices;

namespace StellaOps.TaskRunner.Client.Pagination;

/// <summary>
/// Generic paginator for API responses.
/// </summary>
/// <typeparam name="T">Type of items being paginated.</typeparam>
public sealed class Paginator<T>
{
    private readonly Func<int, int, CancellationToken, Task<PagedResponse<T>>> _fetchPage;
    private readonly int _pageSize;

    /// <summary>
    /// Initializes a new paginator.
    /// </summary>
    /// <param name="fetchPage">Function to fetch a page (offset, limit, cancellationToken) -> page.</param>
    /// <param name="pageSize">Number of items per page (default: 50).</param>
    public Paginator(
        Func<int, int, CancellationToken, Task<PagedResponse<T>>> fetchPage,
        int pageSize = 50)
    {
        _fetchPage = fetchPage ?? throw new ArgumentNullException(nameof(fetchPage));
        _pageSize = pageSize > 0 ? pageSize : throw new ArgumentOutOfRangeException(nameof(pageSize));
    }

    /// <summary>
    /// Iterates through all pages asynchronously.
    /// </summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Async enumerable of items.</returns>
    public async IAsyncEnumerable<T> GetAllAsync([EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        var offset = 0;

        while (true)
        {
            var page = await _fetchPage(offset, _pageSize, cancellationToken).ConfigureAwait(false);

            foreach (var item in page.Items)
            {
                yield return item;
            }

            if (!page.HasMore || page.Items.Count == 0)
            {
                break;
            }

            offset += page.Items.Count;
        }
    }

    /// <summary>
    /// Collects all items into a list.
    /// </summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>List of all items.</returns>
    public async Task<IReadOnlyList<T>> CollectAsync(CancellationToken cancellationToken = default)
    {
        var items = new List<T>();

        await foreach (var item in GetAllAsync(cancellationToken).ConfigureAwait(false))
        {
            items.Add(item);
        }

        return items;
    }

    /// <summary>
    /// Gets a single page.
    /// </summary>
    /// <param name="pageNumber">Page number (1-based).</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Single page response.</returns>
    public Task<PagedResponse<T>> GetPageAsync(int pageNumber, CancellationToken cancellationToken = default)
    {
        if (pageNumber < 1)
        {
            throw new ArgumentOutOfRangeException(nameof(pageNumber), "Page number must be >= 1.");
        }

        var offset = (pageNumber - 1) * _pageSize;
        return _fetchPage(offset, _pageSize, cancellationToken);
    }
}

/// <summary>
/// Paginated response wrapper.
/// </summary>
/// <typeparam name="T">Type of items.</typeparam>
public sealed record PagedResponse<T>(
    IReadOnlyList<T> Items,
    int TotalCount,
    bool HasMore)
{
    /// <summary>
    /// Creates an empty page.
    /// </summary>
    public static PagedResponse<T> Empty { get; } = new([], 0, false);

    /// <summary>
    /// Current page number (1-based) based on offset and page size.
    /// </summary>
    public int PageNumber(int offset, int pageSize)
        => pageSize > 0 ? (offset / pageSize) + 1 : 1;
}

/// <summary>
/// Extension methods for creating paginators.
/// </summary>
public static class PaginatorExtensions
{
    /// <summary>
    /// Creates a paginator from a fetch function.
    /// </summary>
    public static Paginator<T> Paginate<T>(
        this Func<int, int, CancellationToken, Task<PagedResponse<T>>> fetchPage,
        int pageSize = 50)
        => new(fetchPage, pageSize);

    /// <summary>
    /// Takes the first N items from an async enumerable.
    /// </summary>
    public static async IAsyncEnumerable<T> TakeAsync<T>(
        this IAsyncEnumerable<T> source,
        int count,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(source);

        if (count <= 0)
        {
            yield break;
        }

        var taken = 0;
        await foreach (var item in source.WithCancellation(cancellationToken).ConfigureAwait(false))
        {
            yield return item;
            taken++;
            if (taken >= count)
            {
                break;
            }
        }
    }

    /// <summary>
    /// Skips the first N items from an async enumerable.
    /// </summary>
    public static async IAsyncEnumerable<T> SkipAsync<T>(
        this IAsyncEnumerable<T> source,
        int count,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(source);

        var skipped = 0;
        await foreach (var item in source.WithCancellation(cancellationToken).ConfigureAwait(false))
        {
            if (skipped < count)
            {
                skipped++;
                continue;
            }
            yield return item;
        }
    }
}