using System.Runtime.CompilerServices;
using System.Threading.Channels;
using ZB.MOM.WW.LmxProxy.Client.Domain;

namespace ZB.MOM.WW.LmxProxy.Client;

/// <summary>
/// Extension methods for streaming reads, writes, and subscriptions over ILmxProxyClient.
/// </summary>
public static class StreamingExtensions
{
    /// <summary>
    /// Reads multiple tags as an async stream in batches.
    /// Retries up to 2 times per batch. Aborts after 3 consecutive batch errors.
    /// </summary>
    public static async IAsyncEnumerable<KeyValuePair<string, Vtq>> ReadStreamAsync(
        this ILmxProxyClient client,
        IEnumerable<string> addresses,
        int batchSize = 100,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(client);
        ArgumentNullException.ThrowIfNull(addresses);
        if (batchSize <= 0)
            throw new ArgumentOutOfRangeException(nameof(batchSize));

        var batch = new List<string>(batchSize);
        int consecutiveErrors = 0;
        const int maxConsecutiveErrors = 3;
        const int maxRetries = 2;

        foreach (string address in addresses)
        {
            cancellationToken.ThrowIfCancellationRequested();
            batch.Add(address);

            if (batch.Count >= batchSize)
            {
                bool success = false;
                await foreach (var kvp in ReadBatchWithRetry(
                    client, batch, maxRetries, cancellationToken))
                {
                    consecutiveErrors = 0;
                    success = true;
                    yield return kvp;
                }
                if (!success)
                {
                    consecutiveErrors++;
                    if (consecutiveErrors >= maxConsecutiveErrors)
                        yield break;
                }
                batch.Clear();
            }
        }

        // Process remaining
        if (batch.Count > 0)
        {
            await foreach (var kvp in ReadBatchWithRetry(
                client, batch, maxRetries, cancellationToken))
            {
                yield return kvp;
            }
        }
    }

    private static async IAsyncEnumerable<KeyValuePair<string, Vtq>> ReadBatchWithRetry(
        ILmxProxyClient client,
        List<string> batch,
        int maxRetries,
        [EnumeratorCancellation] CancellationToken ct)
    {
        int retries = 0;
        while (retries <= maxRetries)
        {
            IDictionary<string, Vtq>? results = null;
            try
            {
                results = await client.ReadBatchAsync(batch, ct);
            }
            catch when (retries < maxRetries)
            {
                retries++;
                continue;
            }

            if (results is not null)
            {
                foreach (var kvp in results)
                    yield return kvp;
                yield break;
            }
            retries++;
        }
    }

    /// <summary>
    /// Writes values from an async enumerable in batches. Returns total count written.
    /// </summary>
    public static async Task<int> WriteStreamAsync(
        this ILmxProxyClient client,
        IAsyncEnumerable<KeyValuePair<string, TypedValue>> values,
        int batchSize = 100,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(client);
        ArgumentNullException.ThrowIfNull(values);
        if (batchSize <= 0)
            throw new ArgumentOutOfRangeException(nameof(batchSize));

        var batch = new Dictionary<string, TypedValue>(batchSize);
        int totalWritten = 0;

        await foreach (var kvp in values.WithCancellation(cancellationToken))
        {
            batch[kvp.Key] = kvp.Value;

            if (batch.Count >= batchSize)
            {
                await client.WriteBatchAsync(batch, cancellationToken);
                totalWritten += batch.Count;
                batch.Clear();
            }
        }

        if (batch.Count > 0)
        {
            await client.WriteBatchAsync(batch, cancellationToken);
            totalWritten += batch.Count;
        }

        return totalWritten;
    }

    /// <summary>
    /// Processes items in parallel with a configurable max concurrency (default 4).
    /// </summary>
    public static async Task ProcessInParallelAsync<T>(
        this IAsyncEnumerable<T> source,
        Func<T, CancellationToken, Task> processor,
        int maxConcurrency = 4,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(source);
        ArgumentNullException.ThrowIfNull(processor);
        if (maxConcurrency <= 0)
            throw new ArgumentOutOfRangeException(nameof(maxConcurrency));

        using var semaphore = new SemaphoreSlim(maxConcurrency);
        var tasks = new List<Task>();

        await foreach (T item in source.WithCancellation(cancellationToken))
        {
            await semaphore.WaitAsync(cancellationToken);

            tasks.Add(Task.Run(async () =>
            {
                try
                {
                    await processor(item, cancellationToken);
                }
                finally
                {
                    semaphore.Release();
                }
            }, cancellationToken));
        }

        await Task.WhenAll(tasks);
    }

    /// <summary>
    /// Wraps a callback-based subscription into an IAsyncEnumerable via System.Threading.Channels.
    /// </summary>
    public static async IAsyncEnumerable<(string Tag, Vtq Vtq)> SubscribeStreamAsync(
        this ILmxProxyClient client,
        IEnumerable<string> addresses,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(client);
        ArgumentNullException.ThrowIfNull(addresses);

        var channel = Channel.CreateBounded<(string, Vtq)>(
            new BoundedChannelOptions(1000)
            {
                FullMode = BoundedChannelFullMode.DropOldest,
                SingleReader = true,
                SingleWriter = false
            });

        LmxProxyClient.ISubscription? subscription = null;
        try
        {
            subscription = await client.SubscribeAsync(
                addresses,
                (tag, vtq) =>
                {
                    channel.Writer.TryWrite((tag, vtq));
                },
                ex =>
                {
                    channel.Writer.TryComplete(ex);
                },
                cancellationToken);

            await foreach (var item in channel.Reader.ReadAllAsync(cancellationToken))
            {
                yield return item;
            }
        }
        finally
        {
            subscription?.Dispose();
            channel.Writer.TryComplete();
        }
    }
}
