﻿// Adapted from: https://www.codeproject.com/Articles/18409/Asynchronously-Execute-PowerShell-Scripts-from-C

using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.ComponentModel;
using System.Management.Automation;
using System.Management.Automation.Runspaces;
using System.Threading;
using System.Windows.Threading;

namespace AsyncPowerShellScripter
{
    /// <summary>
    /// Class that assists in asynchronously executing and retrieving the results of a powershell script pipeline.
    /// </summary>
    public class PipelineExecutor
    {
        #region public types, members
        public delegate void DataEndDelegate(PipelineExecutor sender);
        public delegate void DataReadyDelegate(PipelineExecutor sender, ICollection<PSObject> data);
        public delegate void ErrorReadyDelegate(PipelineExecutor sender, ICollection<object> data);
        /// <summary>
        /// Occurs when powershell script completed its execution.
        /// </summary>
        public event DataEndDelegate OnDataEnd;
        /// <summary>
        /// Occurs when there is new data available from the powershell script.
        /// </summary>
        public event DataReadyDelegate OnDataReady;
        /// <summary>
        /// Occurs when there is new errordata available from the powershell script.
        /// </summary>
        public event ErrorReadyDelegate OnErrorReady;
        /// <summary>
        /// Gets the powershell Pipeline associated with this PipelineExecutor
        /// </summary>
        public Pipeline Pipeline
        {
            get
            {
                return pipeline;
            }
        }
        #endregion public types, members

        #region private types, members

        /// <summary>
        /// The object that is used to invoke the events on.
        /// </summary>
        private Dispatcher dispatcher;

        /// <summary>
        /// The powershell script pipeline that will be executed asynchronously.
        /// </summary>
        private Pipeline pipeline;

        /// <summary>
        /// Event set when the user wants to stop script execution.
        /// </summary>
        private ManualResetEvent stopEvent;
        /// <summary>
        /// Local delegate to a private method
        /// </summary>
        private DataEndDelegate synchDataEnd;
        /// <summary>
        /// Local delegate to a private method
        /// </summary>
        private DataReadyDelegate synchDataReady;
        /// <summary>
        /// Local delegate to a private method
        /// </summary>
        private ErrorReadyDelegate synchErrorReady;
        /// <summary>
        /// Array of waithandles, used in the StoppableInvoke() method.
        /// </summary>
        private WaitHandle[] waitHandles;
        #endregion private types, members

        #region public methods
        /// <summary>
        /// Constructor, creates a new PipelineExecutor for the given powershell script.
        /// </summary>
        /// <param name="runSpace">Powershell runspace to use for creating and executing the script.</param>
        /// <param name="dispatcher">The object to synchronize the DataReady and DataEnd events with.
        /// Normally you'd pass the form or component here.</param>
        /// <param name="command">The script to run</param>
        public PipelineExecutor(Runspace runSpace, Dispatcher dispatcher, string command)
        {
            this.dispatcher = dispatcher;

            // initialize delegates
            synchDataReady = new DataReadyDelegate(SynchDataReady);
            synchDataEnd = new DataEndDelegate(SynchDataEnd);
            synchErrorReady = new ErrorReadyDelegate(SynchErrorReady);

            // initialize event members
            stopEvent = new ManualResetEvent(false);
            waitHandles = new WaitHandle[] { null, stopEvent };
            // create a pipeline and feed it the script text
            pipeline = runSpace.CreatePipeline(command);

            // we'll listen for script output data by way of the DataReady event
            pipeline.Output.DataReady += new EventHandler(Output_DataReady);
            pipeline.Error.DataReady += new EventHandler(Error_DataReady);
        }

        /// <summary>
        /// Start executing the script in the background.
        /// </summary>
        public void Start()
        {
            if (pipeline.PipelineStateInfo.State == PipelineState.NotStarted)
            {
                // close the pipeline input. If you forget to do
                // this it won't start processing the script.
                pipeline.Input.Close();
                // invoke the pipeline. This will cause it to process the script in the background.
                pipeline.InvokeAsync();
            }
        }
        /// <summary>
        /// Stop executing the script.
        /// </summary>
        public void Stop()
        {
            // first make sure StoppableInvoke() stops blocking
            stopEvent.Set();
            // then tell the pipeline to stop the script
            pipeline.Stop();
        }
        private void Error_DataReady(object sender, EventArgs e)
        {
            // fetch all available objects
            Collection<object> data = pipeline.Error.NonBlockingRead();

            // if there were any, invoke the ErrorReady event
            if (data.Count > 0)
            {
                StoppableInvoke(synchErrorReady, new object[] { this, data });
            }
        }
        #endregion public methods

        #region private methods

        public bool InvokeRequired
        {
            get { return dispatcher.CheckAccess(); }
        }
        public IAsyncResult BeginInvoke(Delegate method, object[] args)
        {
            return new DispatcherOperationAsync(dispatcher.BeginInvoke(method, args));
        }
        public object EndInvoke(IAsyncResult result)
        {
            result.AsyncWaitHandle.WaitOne();
            if (result is DispatcherOperationAsync)
                return ((DispatcherOperationAsync)result).Result;
            return null;
        }
        public object Invoke(Delegate method, object[] args)
        {
            return InvokeRequired ? EndInvoke(BeginInvoke(method, args)) : method.DynamicInvoke(args);
        }
        /// <summary>
        /// Event handler for the DataReady event of the powershell script pipeline.
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void Output_DataReady(object sender, EventArgs e)
        {
            // fetch all available objects
            Collection<PSObject> data = pipeline.Output.NonBlockingRead();

            // if there were any, invoke the DataReady event
            if (data.Count > 0)
            {
                StoppableInvoke(synchDataReady, new object[] { this, data });
            }

            if (pipeline.Output.EndOfPipeline)
            {
                // we're done! invoke the DataEnd event
                StoppableInvoke(synchDataEnd, new object[] { this });
            }
        }
        /// <summary>
        /// Special Invoke method that operates similarly to <see cref="ISynchronizeInvoke.Invoke"/> but in addition to that, it can be
        /// aborted by setting the stopEvent. This avoids potential deadlocks that are possible when using the regular
        /// <see cref="ISynchronizeInvoke.Invoke"/> method.
        /// </summary>
        /// <param name="method">A <see cref="System.Delegate"/> to a method that takes parameters of the same number and type that are
        /// contained in <paramref name="args"/></param>
        /// <param name="args">An array of type <see cref="System.Object"/> to pass as arguments to the given method. This can be null if
        /// no arguments are needed </param>
        /// <returns>The <see cref="Object"/> returned by the asynchronous operation</returns>
        private object StoppableInvoke(Delegate method, object[] args)
        {
            IAsyncResult asyncResult = BeginInvoke(method, args);
            waitHandles[0] = asyncResult.AsyncWaitHandle;
            return (WaitHandle.WaitAny(waitHandles) == 0) ? EndInvoke(asyncResult) : null;
        }
        /// <summary>
        /// private DataEnd handling method that will pass the call on to any handlers that are
        /// attached to the OnDataEnd event of this <see cref="PipelineExecutor"/> instance.
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="data"></param>
        private void SynchDataEnd(PipelineExecutor sender)
        {
            DataEndDelegate delegateDataEndCopy = OnDataEnd;
            if (delegateDataEndCopy != null)
            {
                delegateDataEndCopy(sender);
            }
        }
        /// <summary>
        /// private DataReady handling method that will pass the call on to any event handlers that are
        /// attached to the OnDataReady event of this <see cref="PipelineExecutor"/> instance.
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="data"></param>
        private void SynchDataReady(PipelineExecutor sender, ICollection<PSObject> data)
        {
            DataReadyDelegate delegateDataReadyCopy = OnDataReady;
            if (delegateDataReadyCopy != null)
            {
                delegateDataReadyCopy(sender, data);
            }
        }
        /// <summary>
        /// private ErrorReady handling method that will pass the call on to any event handlers that are
        /// attached to the OnErrorReady event of this <see cref="PipelineExecutor"/> instance.
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="data"></param>
        private void SynchErrorReady(PipelineExecutor sender, ICollection<object> data)
        {
            ErrorReadyDelegate delegateErrorReadyCopy = OnErrorReady;
            if (delegateErrorReadyCopy != null)
            {
                delegateErrorReadyCopy(sender, data);
            }
        }

        #endregion private methods

        #region Internal IAsync Class
        private class DispatcherOperationAsync : IAsyncResult, IDisposable
        {
            private readonly DispatcherOperation _dop;
            private ManualResetEvent _handle = new ManualResetEvent(false);

            #region Implementation of IAsyncResult

            public DispatcherOperationAsync(DispatcherOperation dispatcherOperation)
            {
                _dop = dispatcherOperation;
                _dop.Aborted += DopAborted;
                _dop.Completed += DopCompleted;
            }
            public object AsyncState
            {
                get
                {
                    //Not Implementted
                    return null;
                }
            }
            public WaitHandle AsyncWaitHandle
            {
                get { return _handle; }
            }
            public bool CompletedSynchronously
            {
                get { return false; }
            }
            public bool IsCompleted
            {
                get { return _dop.Status == DispatcherOperationStatus.Completed; }
            }
            public object Result
            {
                get
                {
                    if (!IsCompleted)
                        throw new InvalidAsynchronousStateException("Not Completed");
                    return _dop.Result;
                }
            }
            private void DopAborted(object sender, EventArgs e)
            {
                _handle.Set();
            }
            private void DopCompleted(object sender, EventArgs e)
            {
                _handle.Set();
            }
            #endregion Implementation of IAsyncResult

            #region Implementation of IDisposable

            public void Dispose()
            {
                if (_handle == null) return;
                _handle.Dispose();
                _handle = null;
            }

            #endregion Implementation of IDisposable
        }
        #endregion Internal IAsync Class
    }
}