// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;

#nullable disable

namespace Microsoft.Build.Graph
{
    /// <summary>
    /// Provides deduping of expensive work by a key, or modeling of a set of deduped work that
    /// can be awaited as a unit. Completed results are kept in the collection for reuse.
    /// </summary>
    internal class ParallelWorkSet<TKey, TResult>
    {
        private readonly CancellationToken _cancellationToken;

        /// <summary>
        /// Number of workers to process work items.
        /// </summary>
        private readonly int _degreeOfParallelism;

        private readonly ConcurrentDictionary<TKey, WorkItem> _inProgressOrCompletedWork;

        private bool _isSchedulingCompleted;

        private long _pendingCount;

        private readonly ConcurrentQueue<WorkItem> _queue =
            new ConcurrentQueue<WorkItem>();

        private readonly SemaphoreSlim _semaphore;

        private readonly List<Task> _tasks;

        private readonly List<Exception> _exceptions = new List<Exception>(0);

        /// <summary>
        /// Retrieves all completed work items.
        /// </summary>
        internal Dictionary<TKey, TResult> CompletedWork
        {
            get
            {
                var completedWork = new Dictionary<TKey, TResult>(_inProgressOrCompletedWork.Count);

                foreach (KeyValuePair<TKey, WorkItem> kvp in _inProgressOrCompletedWork)
                {
                    WorkItem workItem = kvp.Value;

                    if (workItem.IsValueCreated)
                    {
                        completedWork[kvp.Key] = workItem.Value;
                    }
                }

                if (_exceptions.Count > 0)
                {
                    throw new AggregateException(_exceptions);
                }

                return completedWork;
            }
        }

        /// <summary>
        /// Checks if the work set has been marked as completed.
        /// </summary>
        internal bool IsCompleted
        {
            get => Volatile.Read(ref _isSchedulingCompleted);
            private set => Volatile.Write(ref _isSchedulingCompleted, value);
        }

        internal ParallelWorkSet(int degreeOfParallelism, IEqualityComparer<TKey> comparer, CancellationToken cancellationToken)
        {
            if (degreeOfParallelism < 0)
            {
                throw new ArgumentException("Degree of parallelism must be a positive integer.");
            }

            _cancellationToken = cancellationToken;
            _degreeOfParallelism = degreeOfParallelism;
            _inProgressOrCompletedWork = new ConcurrentDictionary<TKey, WorkItem>(comparer);

            // Semaphore count is 0 to ensure that all the tasks are blocked unless new data is scheduled.
            _semaphore = new SemaphoreSlim(0, int.MaxValue);
            _tasks = new List<Task>(degreeOfParallelism);

            for (int i = 0; i < degreeOfParallelism; i++)
            {
                _tasks.Add(CreateProcessorItemTask());
            }
        }

        /// <summary>
        /// Enqueues a work item to the work set.
        /// </summary>
        /// <param name="key"></param>
        /// <param name="workFunc"></param>
        internal void AddWork(TKey key, Func<TResult> workFunc)
        {
            if (IsCompleted)
            {
                throw new InvalidOperationException("Cannot add new work after work set is marked as completed.");
            }

            WorkItem workItem = new(workFunc);

            if (!_inProgressOrCompletedWork.TryAdd(key, workItem))
            {
                return;
            }

            Interlocked.Increment(ref _pendingCount);

            // NOTE: Enqueue MUST happen before releasing the semaphore
            // to ensure WaitAsync below never returns when there is not
            // a corresponding item in the queue to be dequeued. The only
            // exception is on completion of all items.
            _queue.Enqueue(workItem);
            _semaphore.Release();
        }

        /// <summary>
        /// Assists processing items until all the items added to the queue are processed, completes the work set, and
        /// propagates any exceptions thrown by workers.
        /// </summary>
        internal void WaitForAllWorkAndComplete()
        {
            if (IsCompleted)
            {
                return;
            }
            while (!_cancellationToken.IsCancellationRequested && Interlocked.Read(ref _pendingCount) > 0)
            {
                ExecuteWorkItem();
            }

            IsCompleted = true;

            // Release one thread that will release all the threads when all the elements are processed.
            _semaphore.Release();
            Task.WaitAll(
#if NET
                _tasks);
#else
                _tasks.ToArray());
#endif

            if (_exceptions.Count > 0)
            {
                throw new AggregateException(_exceptions);
            }
        }

        private Task CreateProcessorItemTask()
        {
            return Task.Run(
                async () =>
                {
                    bool shouldStopAllWorkers = false;

                    while (!shouldStopAllWorkers)
                    {
                        await _semaphore.WaitAsync(_cancellationToken);

                        try
                        {
                            ExecuteWorkItem();
                        }
                        finally
                        {
                            shouldStopAllWorkers = Interlocked.Read(ref _pendingCount) == 0 && IsCompleted;

                            if (shouldStopAllWorkers)
                            {
                                // Ensure all tasks are unblocked and can gracefully
                                // finish since there are at most degreeOfParallelism - 1 tasks
                                // waiting at this point
                                _semaphore.Release(_degreeOfParallelism);
                            }
                        }
                    }
                },
                _cancellationToken);
        }

        private void ExecuteWorkItem()
        {
            if (_queue.TryDequeue(out WorkItem workItem))
            {
                try
                {
                    _ = workItem.Value;
                }
                catch (Exception ex)
                {
                    lock (_exceptions)
                    {
                        _exceptions.Add(ex);
                    }
                }
                finally
                {
                    Interlocked.Decrement(ref _pendingCount);
                }
            }
        }

        /// <summary>
        /// A lazily-evaluated work item. Unlike <see cref="Lazy{T}"/>, this does not place a
        /// <c>DynamicallyAccessedMembers</c> requirement on <typeparamref name="TResult"/> (<see cref="Lazy{T}"/>
        /// annotates its type argument for the parameterless constructor, which triggers IL2091 here), and it is
        /// only ever evaluated once by a single worker before the result is read after
        /// <see cref="WaitForAllWorkAndComplete"/> establishes a happens-before barrier.
        /// </summary>
        private sealed class WorkItem
        {
            private Func<TResult> _workFunc;
            private bool _isValueCreated;
            private TResult _value;

            internal WorkItem(Func<TResult> workFunc)
            {
                _workFunc = workFunc;
            }

            internal bool IsValueCreated => _isValueCreated;

            internal TResult Value
            {
                get
                {
                    if (!_isValueCreated)
                    {
                        _value = _workFunc();
                        _isValueCreated = true;

                        // Release the factory (and anything its closure captured) once the value is
                        // computed. Completed work items are retained in the set for reuse, so holding
                        // the delegate would keep captured graph state alive for the whole build.
                        // Matches Lazy<T>, which also drops its value factory after evaluation.
                        _workFunc = null;
                    }

                    return _value;
                }
            }
        }
    }
}