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

namespace ZB.MOM.WW.MxGateway.Worker.Sta;

/// <summary>
/// Encapsulates a work item to be executed on an STA thread with cancellation support.
/// </summary>
internal sealed class StaWorkItem<T> : IStaWorkItem
{
    private readonly Func<T> command;
    private readonly CancellationToken cancellationToken;
    private readonly CancellationTokenRegistration cancellationRegistration;
    private int started;

    /// <summary>Initializes a work item with a command and cancellation token.</summary>
    /// <param name="command">Function to execute on the STA thread.</param>
    /// <param name="cancellationToken">Token to cancel the work item.</param>
    public StaWorkItem(Func<T> command, CancellationToken cancellationToken)
    {
        this.command = command ?? throw new ArgumentNullException(nameof(command));
        this.cancellationToken = cancellationToken;
        Completion = new TaskCompletionSource<T>(TaskCreationOptions.RunContinuationsAsynchronously);

        if (cancellationToken.CanBeCanceled)
        {
            cancellationRegistration = cancellationToken.Register(
                () =>
                {
                    if (Interlocked.CompareExchange(ref started, 1, 0) == 0)
                    {
                        Completion.TrySetCanceled(cancellationToken);
                    }
                });
        }
    }

    /// <summary>Gets the task that completes when work completes.</summary>
    public Task<T> Task => Completion.Task;

    private TaskCompletionSource<T> Completion { get; }

    /// <inheritdoc />
    public void CancelBeforeExecution()
    {
        if (Interlocked.CompareExchange(ref started, 1, 0) == 0)
        {
            Completion.TrySetCanceled(cancellationToken);
            cancellationRegistration.Dispose();
        }
    }

    /// <inheritdoc />
    public void Execute()
    {
        if (Interlocked.CompareExchange(ref started, 1, 0) != 0)
        {
            cancellationRegistration.Dispose();
            return;
        }

        cancellationRegistration.Dispose();

        if (cancellationToken.IsCancellationRequested)
        {
            Completion.TrySetCanceled(cancellationToken);
            return;
        }

        try
        {
            Completion.TrySetResult(command());
        }
        catch (Exception exception)
        {
            Completion.TrySetException(exception);
        }
    }
}
