﻿/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

namespace DotPulsar.Internal;

using Exceptions;
using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;

public sealed class AsyncLock : IAsyncDisposable
{
    private readonly LinkedList<CancelableCompletionSource<IDisposable>> _pending;
    private readonly SemaphoreSlim _semaphoreSlim;
    private readonly Releaser _releaser;
    private readonly Task<IDisposable> _completedTask;
    private int _isDisposed;

    public AsyncLock()
    {
        _pending = new LinkedList<CancelableCompletionSource<IDisposable>>();
        _semaphoreSlim = new SemaphoreSlim(1, 1);
        _releaser = new Releaser(Release);
        _completedTask = Task.FromResult((IDisposable) _releaser);
    }

    public Task<IDisposable> Lock(CancellationToken cancellationToken)
    {
        LinkedListNode<CancelableCompletionSource<IDisposable>>? node = null;

        lock (_pending)
        {
            ThrowIfDisposed();

            if (_semaphoreSlim.CurrentCount == 1) //Lock is free
            {
                _semaphoreSlim.Wait(cancellationToken); //Will never block
                return _completedTask;
            }

            //Lock was not free
            var ccs = new CancelableCompletionSource<IDisposable>();
            node = _pending.AddLast(ccs);
        }

        cancellationToken.Register(() => Cancel(node));

        return node.Value.Task;
    }

    public async ValueTask DisposeAsync()
    {
        lock (_pending)
        {
            if (Interlocked.Exchange(ref _isDisposed, 1) != 0)
                return;

            foreach (var pending in _pending)
                pending.Dispose();

            _pending.Clear();
        }

        await _semaphoreSlim.WaitAsync().ConfigureAwait(false); //Wait for possible lock-holder to finish

        _semaphoreSlim.Release();
        _semaphoreSlim.Dispose();
    }

    private void Cancel(LinkedListNode<CancelableCompletionSource<IDisposable>> node)
    {
        lock (_pending)
        {
            try
            {
                _pending.Remove(node);
                node.Value.Dispose();
            }
            catch
            {
                // Ignore
            }
        }
    }

    private void Release()
    {
        lock (_pending)
        {
            var node = _pending.First;
            if (node is not null)
            {
                node.Value.SetResult(_releaser);
                node.Value.Dispose();
                _pending.RemoveFirst();
                return;
            }

            if (_semaphoreSlim.CurrentCount == 0)
                _semaphoreSlim.Release();
        }
    }

    private void ThrowIfDisposed()
    {
        if (_isDisposed != 0)
            throw new AsyncLockDisposedException();
    }

    private class Releaser : IDisposable
    {
        private readonly Action _release;

        public Releaser(Action release)
            => _release = release;

        public void Dispose()
            => _release();
    }
}