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

namespace MxGateway.Worker.Sta;

internal sealed class StaWorkItem<T> : IStaWorkItem
{
    private readonly Func<T> command;
    private readonly CancellationToken cancellationToken;
    private readonly CancellationTokenRegistration cancellationRegistration;
    private int started;

    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);
                    }
                });
        }
    }

    public Task<T> Task => Completion.Task;

    private TaskCompletionSource<T> Completion { get; }

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

    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);
        }
    }
}