// -----------------------------------------------------------------------------
// <copyright file="PowerShellCmdlet.cs" company="Microsoft Corporation">
//     Copyright (c) Microsoft Corporation. Licensed under the MIT License.
// </copyright>
// -----------------------------------------------------------------------------

namespace Microsoft.WinGet.Common.Command
{
    using System;
    using System.Collections.Concurrent;
    using System.Collections.Generic;
    using System.Management.Automation;
    using System.Runtime.ExceptionServices;
    using System.Threading;
    using System.Threading.Tasks;
    using Microsoft.WinGet.Resources;
    using Microsoft.WinGet.SharedLib.Exceptions;
    using Microsoft.WinGet.SharedLib.PolicySettings;

    /// <summary>
    /// This must be the base class for every cmdlet for winget PowerShell modules.
    /// It supports:
    ///  - Async operations.
    ///  - Execute on an MTA. If the thread is already running on an MTA it will executed it, otherwise
    ///    it will create a new MTA thread.
    /// Wait must be used to synchronously wait con the task.
    /// </summary>
    public abstract class PowerShellCmdlet
    {
        private const string Debug = "Debug";
        private static readonly string[] WriteInformationTags = new string[] { "PSHOST" };

        private readonly PSCmdlet psCmdlet;
        private readonly Thread pwshThread;

        private readonly CancellationTokenSource source = new ();
        private readonly SemaphoreSlim semaphore = new (1, 1);
        private readonly ManualResetEventSlim pwshThreadActionReady = new (false);
        private readonly ManualResetEventSlim pwshThreadActionCompleted = new (false);

        private BlockingCollection<QueuedStream> queuedStreams = new ();
        private int progressActivityId = 0;
        private ConcurrentDictionary<int, ProgressRecordType> progressRecords = new ();
        private Action? pwshThreadAction = null;
        private ExceptionDispatchInfo? pwshThreadEdi = null;

        /// <summary>
        /// Initializes a new instance of the <see cref="PowerShellCmdlet"/> class.
        /// </summary>
        /// <param name="psCmdlet">PSCmdlet.</param>
        /// <param name="policies">Policies.</param>
        public PowerShellCmdlet(PSCmdlet psCmdlet, HashSet<Policy> policies)
        {
            // Passing Debug will make all the message actions to be Inquire. For async operations
            // and the current queue message implementation this doesn't make sense.
            // PowerShell will inquire for any message giving the impression that the task is
            // paused, but the async operation is still running.
            if (psCmdlet.MyInvocation.BoundParameters.ContainsKey(Debug))
            {
                throw new NotSupportedException(Resources.DebugNotSupported);
            }

            this.ValidatePolicies(policies);

            this.psCmdlet = psCmdlet;
            this.pwshThread = Thread.CurrentThread;
        }

        /// <summary>
        /// Request cancellation for this command.
        /// </summary>
        public void Cancel()
        {
            this.source.Cancel();
        }

        /// <summary>
        /// Execute the delegate in a MTA thread.
        /// Caller must wait on task.
        /// </summary>
        /// <param name="func">Function to execute.</param>
        /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
        internal Task RunOnMTA(Func<Task> func)
        {
            // .NET 4.8 doesn't support TaskCompletionSource.
#if POWERSHELL_WINDOWS
            throw new NotImplementedException();
#else
            // This must be called in the main thread.
            if (this.pwshThread != Thread.CurrentThread)
            {
                throw new InvalidOperationException();
            }

            if (Thread.CurrentThread.GetApartmentState() == ApartmentState.MTA)
            {
                this.Write(StreamType.Verbose, "Already running on MTA");
                try
                {
                    Task result = func();
                    result.ContinueWith((task) => this.Complete(), TaskContinuationOptions.ExecuteSynchronously);
                    return result;
                }
                catch
                {
                    this.Complete();
                    throw;
                }
            }

            this.Write(StreamType.Verbose, "Creating MTA thread");
            var tcs = new TaskCompletionSource();
            var thread = new Thread(() =>
            {
                try
                {
                    func().GetAwaiter().GetResult();
                    tcs.SetResult();
                }
                catch (Exception e)
                {
                    tcs.SetException(e);
                }
                finally
                {
                    this.Complete();
                }
            });

            thread.SetApartmentState(ApartmentState.MTA);
            thread.Start();
            return tcs.Task;
#endif
        }

        /// <summary>
        /// Execute the delegate in a MTA thread.
        /// Caller must wait on task.
        /// </summary>
        /// <param name="func">Function to execute.</param>
        /// <typeparam name="TResult">Return type of function.</typeparam>
        /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
        internal Task<TResult> RunOnMTA<TResult>(Func<Task<TResult>> func)
        {
            // This must be called in the main thread.
            if (this.pwshThread != Thread.CurrentThread)
            {
                throw new InvalidOperationException();
            }

            if (Thread.CurrentThread.GetApartmentState() == ApartmentState.MTA)
            {
                this.Write(StreamType.Verbose, "Already running on MTA");
                try
                {
                    Task<TResult> result = func();
                    result.ContinueWith((task) => this.Complete(), TaskContinuationOptions.ExecuteSynchronously);
                    return result;
                }
                catch
                {
                    this.Complete();
                    throw;
                }
            }

            this.Write(StreamType.Verbose, "Creating MTA thread");
            var tcs = new TaskCompletionSource<TResult>();
            var thread = new Thread(() =>
            {
                try
                {
                    var result = func().GetAwaiter().GetResult();
                    tcs.SetResult(result);
                }
                catch (Exception e)
                {
                    tcs.SetException(e);
                }
                finally
                {
                    this.Complete();
                }
            });

            thread.SetApartmentState(ApartmentState.MTA);
            thread.Start();
            return tcs.Task;
        }

        /// <summary>
        /// Execute the delegate in a MTA thread.
        /// Synchronous call.
        /// </summary>
        /// <param name="func">Function to execute.</param>
        /// <typeparam name="TResult">Return type of function.</typeparam>
        /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
        internal TResult RunOnMTA<TResult>(Func<TResult> func)
        {
            // This must be called in the main thread.
            if (this.pwshThread != Thread.CurrentThread)
            {
                throw new InvalidOperationException();
            }

            if (Thread.CurrentThread.GetApartmentState() == ApartmentState.MTA)
            {
                this.Write(StreamType.Verbose, "Already running on MTA");
                try
                {
                    return func();
                }
                finally
                {
                    this.Complete();
                }
            }

            this.Write(StreamType.Verbose, "Creating MTA thread");
            var tcs = new TaskCompletionSource<TResult>();
            var thread = new Thread(() =>
            {
                try
                {
                    var result = func();
                    tcs.SetResult(result);
                }
                catch (Exception e)
                {
                    tcs.SetException(e);
                }
                finally
                {
                    this.Complete();
                }
            });

            thread.SetApartmentState(ApartmentState.MTA);
            thread.Start();
            this.Wait(tcs.Task);
            return tcs.Task.Result;
        }

        /// <summary>
        /// Executes an action in the main thread.
        /// Blocks until call is executed.
        /// </summary>
        /// <param name="action">Action to perform.</param>
        internal void ExecuteInPowerShellThread(Action action)
        {
            if (this.pwshThread == Thread.CurrentThread)
            {
                action();
                return;
            }

            this.WaitForOurTurn();

            this.pwshThreadAction = action;
            this.pwshThreadActionReady.Set();
            this.WaitMainThreadActionCompletion();
        }

        /// <summary>
        /// Waits for the task to be completed. This MUST be called from the main thread.
        /// </summary>
        /// <param name="runningTask">Task to wait for.</param>
        /// <param name="writeCmdlet">The cmdlet that can write to PowerShell.</param>
        internal void Wait(Task runningTask, PowerShellCmdlet? writeCmdlet = null)
        {
            writeCmdlet ??= this;

            // This must be called in the main thread.
            if (this.pwshThread != Thread.CurrentThread)
            {
                throw new InvalidOperationException();
            }

            do
            {
                if (this.pwshThreadActionReady.IsSet)
                {
                    // Someone needs the main thread.
                    this.pwshThreadActionReady.Reset();

                    if (this.pwshThreadAction != null)
                    {
                        try
                        {
                            this.pwshThreadEdi = null;
                            this.pwshThreadAction();
                        }
                        catch (Exception e)
                        {
                            // Make sure we don't throw in the PowerShell thread, this way
                            // we'll get a more meaningful stack by Get-Error.
                            this.pwshThreadEdi = ExceptionDispatchInfo.Capture(e);
                        }

                        this.pwshThreadAction = null;
                    }

                    // Done.
                    this.pwshThreadActionCompleted.Set();
                }

                // Take from the blocking collection.
                if (!this.queuedStreams.IsCompleted && this.queuedStreams.Count > 0)
                {
                    try
                    {
                        var queuedOutput = this.queuedStreams.Take();
                        if (queuedOutput != null)
                        {
                            this.CmdletWrite(queuedOutput.Type, queuedOutput.Data, writeCmdlet);
                        }
                    }
                    catch (InvalidOperationException)
                    {
                        // An InvalidOperationException means that Take() was called on a completed collection.
                    }
                }
            }
            while (!(runningTask.IsCompleted && this.queuedStreams.IsCompleted));

            if (runningTask.IsFaulted)
            {
                // If IsFaulted is true, the task's Status is equal to Faulted,
                // and its Exception property will be non-null.
                AggregateException? ae = runningTask.Exception! as AggregateException;
                if (ae != null && ae.InnerExceptions.Count == 1)
                {
                    ExceptionDispatchInfo.Capture(ae.InnerExceptions[0]).Throw();
                }

                throw runningTask.Exception!;
            }
        }

        /// <summary>
        /// Writes into the corresponding stream if running on the main thread.
        /// Otherwise queue the message.
        /// Is the caller responsibility to use the correct types.
        /// </summary>
        /// <param name="type">Stream type.</param>
        /// <param name="data">Data.</param>
        internal void Write(StreamType type, object data)
        {
            if (type == StreamType.Progress)
            {
                ProgressRecord progressRecord = (ProgressRecord)data;
                if (progressRecord.RecordType == ProgressRecordType.Completed)
                {
                    throw new NotSupportedException("Use CompleteProgress");
                }

                // Keep track of all progress activity.
                _ = this.progressRecords.TryAdd(progressRecord.ActivityId, progressRecord.RecordType);
            }

            if (this.pwshThread == Thread.CurrentThread)
            {
                this.CmdletWrite(type, data, this);
                return;
            }

            this.queuedStreams.Add(new QueuedStream(type, data));
        }

        /// <summary>
        /// Helper to compute percentage and write progress for processing activities.
        /// </summary>
        /// <param name="activityId">Activity id.</param>
        /// <param name="activity">The activity in progress.</param>
        /// <param name="status">The status of the activity.</param>
        /// <param name="completed">Number of completed actions.</param>
        /// <param name="total">The expected total.</param>
        internal void WriteProgressWithPercentage(int activityId, string activity, string status, int completed, int total)
        {
            double percentComplete = (double)completed / total;
            var record = new ProgressRecord(activityId, activity, status)
            {
                RecordType = ProgressRecordType.Processing,
                PercentComplete = (int)(100.0 * percentComplete),
            };
            this.Write(StreamType.Progress, record);
        }

        /// <summary>
        /// Helper to complete progress records.
        /// </summary>
        /// <param name="activityId">Activity id.</param>
        /// <param name="activity">The activity in progress.</param>
        /// <param name="status">The status of the activity.</param>
        /// <param name="force">Force write complete progress.</param>
        internal void CompleteProgress(int activityId, string activity, string status, bool force = false)
        {
            var record = new ProgressRecord(activityId, activity, status)
            {
                RecordType = ProgressRecordType.Completed,
                PercentComplete = 100,
            };

            if (!this.progressRecords.TryAdd(activityId, record.RecordType))
            {
                _ = this.progressRecords.TryUpdate(activityId, record.RecordType, ProgressRecordType.Processing);
            }

            if (this.pwshThread == Thread.CurrentThread)
            {
                this.CmdletWrite(StreamType.Progress, record, this);
            }
            else
            {
                // You should only use force if you know the cmdlet that is completing this progress is a sync cmdlet that
                // is running in an async context. A sync cmdlet is anything that doesn't start with Start-*
                if (force)
                {
                    this.ExecuteInPowerShellThread(() => this.CmdletWrite(StreamType.Progress, record, this));
                }
                else
                {
                    this.queuedStreams.Add(new QueuedStream(StreamType.Progress, record));
                }
            }
        }

        /// <summary>
        /// Writes to PowerShell streams.
        /// This method must be called in the original thread.
        /// WARNING: You must only call this when the task is completed.
        /// </summary>
        /// <param name="writeCmdlet">The cmdlet that can write to PowerShell.</param>
        internal void ConsumeAndWriteStreams(PowerShellCmdlet writeCmdlet)
        {
            // This must be called in the main thread.
            if (this.pwshThread != Thread.CurrentThread)
            {
                throw new InvalidOperationException();
            }

            // Take from the blocking collection until is completed.
            try
            {
                while (true)
                {
                    var queuedOutput = this.queuedStreams.Take();
                    if (queuedOutput != null)
                    {
                        this.CmdletWrite(queuedOutput.Type, queuedOutput.Data, writeCmdlet);
                    }
                }
            }
            catch (InvalidOperationException)
            {
                // We are done.
                // An InvalidOperationException means that Take() was called on a completed collection.
            }
        }

        /// <summary>
        /// Gets a new progress activity id.
        /// </summary>
        /// <returns>The new progress record id.</returns>
        internal int GetNewProgressActivityId()
        {
            return Interlocked.Increment(ref this.progressActivityId);
        }

        /// <summary>
        /// Gets the cancellation token.
        /// </summary>
        /// <returns>CancellationToken.</returns>
        internal CancellationToken GetCancellationToken()
        {
            return this.source.Token;
        }

        /// <summary>
        /// Gets the current file system location from the cmdlet.
        /// </summary>
        /// <returns>Path.</returns>
        internal string GetCurrentFileSystemLocation()
        {
            return this.psCmdlet.SessionState.Path.CurrentFileSystemLocation.Path;
        }

        /// <summary>
        /// Sets a variable.
        /// </summary>
        /// <param name="variableName">Variable name.</param>
        /// <param name="value">Value.</param>
        internal void SetVariable(string variableName, object value)
        {
            this.psCmdlet.SessionState.PSVariable.Set(variableName, value);
        }

        /// <summary>
        /// Prompts the user if it should continue processing if possible.
        /// </summary>
        /// <param name="target">Message.</param>
        /// <returns>If the operation should continue.</returns>
        internal bool ShouldProcess(string target)
        {
            // If not on the main thread just continue.
            if (this.pwshThread != Thread.CurrentThread)
            {
                return true;
            }

            return this.psCmdlet.ShouldProcess(target);
        }

        private void Complete()
        {
            this.queuedStreams.CompleteAdding();
        }

        private void CmdletWrite(StreamType streamType, object data, PowerShellCmdlet writeCmdlet)
        {
            switch (streamType)
            {
                case StreamType.Debug:
                    throw new NotSupportedException();
                case StreamType.Verbose:
                    writeCmdlet.psCmdlet.WriteVerbose((string)data);
                    break;
                case StreamType.Warning:
                    writeCmdlet.psCmdlet.WriteWarning((string)data);
                    break;
                case StreamType.Error:
                    writeCmdlet.psCmdlet.WriteError((ErrorRecord)data);
                    break;
                case StreamType.Progress:
                    // If the activity is already completed don't write progress.
                    var progressRecord = (ProgressRecord)data;
                    if (this.progressRecords[progressRecord.ActivityId] == progressRecord.RecordType)
                    {
                        writeCmdlet.psCmdlet.WriteProgress(progressRecord);
                    }

                    break;
                case StreamType.Object:
                    writeCmdlet.psCmdlet.WriteObject(data);
                    break;
                case StreamType.Information:
                    writeCmdlet.psCmdlet.WriteInformation(data, WriteInformationTags);
                    break;
            }
        }

        private void ValidatePolicies(HashSet<Policy> policies)
        {
            GroupPolicy groupPolicy = GroupPolicy.GetInstance();

            if (policies.Contains(Policy.WinGet))
            {
                if (!groupPolicy.IsEnabled(Policy.WinGet))
                {
                    throw new GroupPolicyException(Policy.WinGet, GroupPolicyFailureType.BlockedByPolicy);
                }

                policies.Remove(Policy.WinGet);
            }

            if (policies.Contains(Policy.Configuration))
            {
                if (!groupPolicy.IsEnabled(Policy.Configuration))
                {
                    throw new GroupPolicyException(Policy.Configuration, GroupPolicyFailureType.BlockedByPolicy);
                }

                policies.Remove(Policy.Configuration);
            }

            if (policies.Contains(Policy.WinGetCommandLineInterfaces))
            {
                if (!groupPolicy.IsEnabled(Policy.WinGetCommandLineInterfaces))
                {
                    throw new GroupPolicyException(Policy.WinGetCommandLineInterfaces, GroupPolicyFailureType.BlockedByPolicy);
                }

                policies.Remove(Policy.WinGetCommandLineInterfaces);
            }

            if (policies.Count > 0)
            {
                throw new NotSupportedException($"Invalid policies {string.Join(",", policies)}");
            }
        }

        private void WaitForOurTurn()
        {
            this.semaphore.Wait(this.GetCancellationToken());
            this.pwshThreadActionCompleted.Reset();
        }

        private void WaitMainThreadActionCompletion()
        {
            WaitHandle.WaitAny(new[]
            {
                this.GetCancellationToken().WaitHandle,
                this.pwshThreadActionCompleted.WaitHandle,
            });

            try
            {
                if (this.pwshThreadEdi != null)
                {
                    this.pwshThreadEdi.Throw();
                }
            }
            finally
            {
                this.semaphore.Release();
            }
        }

        private class QueuedStream
        {
            public QueuedStream(StreamType type, object data)
            {
                this.Type = type;
                this.Data = data;
            }

            public StreamType Type { get; }

            public object Data { get; }
        }
    }
}
