using System;
using System.Threading.Tasks;
using System.Windows.Threading;

namespace Laservall.Solidworks.Server
{
    /// <summary>
    /// Marshals work to the SolidWorks main STA thread via the WPF Dispatcher.
    /// All SolidWorks COM API calls MUST go through this marshaller to avoid
    /// RPC_E_DISCONNECTED / DisconnectedContext errors.
    /// </summary>
    internal sealed class StaThreadMarshaller
    {
        private readonly Dispatcher _dispatcher;

        public StaThreadMarshaller(Dispatcher dispatcher)
        {
            _dispatcher = dispatcher ?? throw new ArgumentNullException(nameof(dispatcher));
        }

        /// <summary>
        /// Run an action synchronously on the STA thread. Blocks until complete.
        /// Use from background threads when you need the result immediately.
        /// </summary>
        public void Invoke(Action action)
        {
            if (_dispatcher.CheckAccess())
            {
                action();
                return;
            }

            _dispatcher.Invoke(action);
        }

        /// <summary>
        /// Run a function synchronously on the STA thread and return its result.
        /// Blocks the calling thread until complete.
        /// </summary>
        public T Invoke<T>(Func<T> func)
        {
            if (_dispatcher.CheckAccess())
            {
                return func();
            }

            return _dispatcher.Invoke(func);
        }

        /// <summary>
        /// Run a function asynchronously on the STA thread and return a Task with the result.
        /// Does not block the calling thread.
        /// </summary>
        public async Task<T> InvokeAsync<T>(Func<T> func)
        {
            if (_dispatcher.CheckAccess())
            {
                return func();
            }

            return await _dispatcher.InvokeAsync(func);
        }

        /// <summary>
        /// Run an async function on the STA thread. The async continuations will
        /// remain on the STA thread (DispatcherSynchronizationContext), keeping
        /// all COM calls safe throughout the entire async chain.
        /// </summary>
        public Task<T> InvokeAsync<T>(Func<Task<T>> asyncFunc)
        {
            if (_dispatcher.CheckAccess())
            {
                return asyncFunc();
            }

            var tcs = new TaskCompletionSource<T>();
            _dispatcher.InvokeAsync(async () =>
            {
                try
                {
                    T result = await asyncFunc();
                    tcs.SetResult(result);
                }
                catch (Exception ex)
                {
                    tcs.SetException(ex);
                }
            });
            return tcs.Task;
        }

        /// <summary>
        /// Run an async action on the STA thread (no return value).
        /// </summary>
        public Task InvokeAsync(Func<Task> asyncFunc)
        {
            if (_dispatcher.CheckAccess())
            {
                return asyncFunc();
            }

            var tcs = new TaskCompletionSource<bool>();
            _dispatcher.InvokeAsync(async () =>
            {
                try
                {
                    await asyncFunc();
                    tcs.SetResult(true);
                }
                catch (Exception ex)
                {
                    tcs.SetException(ex);
                }
            });
            return tcs.Task;
        }
    }
}