// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.
 
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.VisualStudio.Extensibility.Testing;
using Microsoft.VisualStudio.Language.Intellisense;
using Microsoft.VisualStudio.Text.Editor;
using Microsoft.VisualStudio.Threading;
 
namespace Roslyn.VisualStudio.IntegrationTests.InProcess
{
    internal static class LightBulbHelper
    {
        public static async Task<bool> WaitForLightBulbSessionAsync(TestServices testServices, ILightBulbBroker broker, IWpfTextView view, CancellationToken cancellationToken)
        {
            await testServices.Editor.WaitForEditorOperationsAsync(cancellationToken);
            var active = broker.IsLightBulbSessionActive(view);
            if (!active)
                return false;
 
            await WaitForItemsAsync(testServices, broker, view, cancellationToken);
            return true;
        }
 
        public static async Task<IEnumerable<SuggestedActionSet>> WaitForItemsAsync(TestServices testServices, ILightBulbBroker broker, IWpfTextView view, CancellationToken cancellationToken)
        {
            while (true)
            {
                var items = await TryWaitForItemsAsync(testServices, broker, view, cancellationToken);
                if (items is not null)
                    return items;
 
                // The session was dismissed unexpectedly. The editor might show it again.
                await testServices.Editor.WaitForEditorOperationsAsync(cancellationToken);
            }
        }
 
        private static async Task<IEnumerable<SuggestedActionSet>?> TryWaitForItemsAsync(TestServices testServices, ILightBulbBroker broker, IWpfTextView view, CancellationToken cancellationToken)
        {
            var activeSession = broker.GetSession(view);
            if (activeSession == null)
            {
                var bufferType = view.TextBuffer.ContentType.DisplayName;
                throw new InvalidOperationException($"No expanded light bulb session found after View.ShowSmartTag.  Buffer content type={bufferType}");
            }
 
            var asyncSession = (IAsyncLightBulbSession)activeSession;
            var tcs = new TaskCompletionSource<List<SuggestedActionSet>>();
 
            EventHandler<SuggestedActionsUpdatedArgs>? handler = null;
            handler = (s, e) =>
            {
                // ignore these.  we care about when the lightbulb items are all completed.
                if (e.Status == QuerySuggestedActionCompletionStatus.InProgress)
                    return;
 
                if (e.Status == QuerySuggestedActionCompletionStatus.Completed)
                    tcs.SetResult(e.ActionSets.ToList());
                else if (e.Status == QuerySuggestedActionCompletionStatus.CompletedWithoutData)
                    tcs.SetResult(new List<SuggestedActionSet>());
                else if (e.Status == QuerySuggestedActionCompletionStatus.Canceled)
                    tcs.TrySetCanceled();
                else
                    tcs.TrySetException(new InvalidOperationException($"Light bulb transitioned to non-complete state: {e.Status}"));
 
                asyncSession.SuggestedActionsUpdated -= handler;
            };
 
            asyncSession.SuggestedActionsUpdated += handler;
 
            asyncSession.Dismissed += (_, _) => tcs.TrySetCanceled(new CancellationToken(true));
 
            if (asyncSession.IsDismissed)
                tcs.TrySetCanceled(new CancellationToken(true));
 
            // Calling PopulateWithData ensures the underlying session will call SuggestedActionsUpdated at least once
            // with the latest data computed.  This is needed so that if the lightbulb computation is already complete
            // that we hear about the results.
            await asyncSession.PopulateWithDataAsync(overrideRequestedActionCategories: null, operationContext: null).ConfigureAwait(false);
 
            try
            {
                return await tcs.Task.WithCancellation(cancellationToken);
            }
            catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
            {
                var version = await testServices.Shell.GetVersionAsync(cancellationToken);
                if (Version.Parse("17.2.32427.441") >= version)
                {
                    // Unexpected cancellation can occur when the editor dismisses the light bulb without request
                    return null;
                }
 
                throw new OperationCanceledException($"IDE version '{version}' unexpectedly dismissed the light bulb.");
            }
        }
    }
}