#nullable enable

using System;
using System.Collections.Concurrent;
using System.Diagnostics;
using System.Threading;
using Cysharp.Threading.Tasks;

namespace ShrinkEventBus
{
    internal interface IShrinkBusScheduler : IDisposable
    {
        bool IsOnSchedulerThread { get; }
        bool TryPost(Action action);
        UniTask PostAsync(Func<UniTask> action, CancellationToken cancellationToken);
        UniTask ShutdownAsync(bool drain, TimeSpan timeout);
    }

    internal sealed class ShrinkQueueFullException : InvalidOperationException
    {
        public ShrinkQueueFullException(string schedulerName)
            : base($"Scheduler queue '{schedulerName}' is full.")
        {
        }
    }

    internal static class ShrinkBusSchedulerFactory
    {
        public static IShrinkBusScheduler Create(ShrinkBusOptions options, string name)
        {
            return options.Scheduler switch
            {
                ShrinkBusSchedulerKind.MainThread => new ShrinkMainThreadScheduler(name, options),
                ShrinkBusSchedulerKind.DedicatedThread => new ShrinkDedicatedThreadScheduler(name, options),
                ShrinkBusSchedulerKind.TaskPool => new ShrinkTaskPoolScheduler(name, options),
                _ => new ShrinkInlineScheduler()
            };
        }
    }

    internal sealed class ShrinkSchedulerWorkItem
    {
        private readonly UniTaskCompletionSource? _completion;

        public ShrinkSchedulerWorkItem(Func<UniTask> action, bool awaitable)
        {
            Action = action ?? throw new ArgumentNullException(nameof(action));
            if (awaitable)
                _completion = new UniTaskCompletionSource();
        }

        public Func<UniTask> Action { get; }
        public UniTask Completion => _completion?.Task ?? UniTask.CompletedTask;

        public void Complete() => _completion?.TrySetResult();
        public void Fail(Exception exception) => _completion?.TrySetException(exception);
        public void Cancel(CancellationToken cancellationToken) => _completion?.TrySetCanceled(cancellationToken);
        public void Drop(string schedulerName) =>
            _completion?.TrySetException(new ShrinkQueueFullException(schedulerName));
    }

    internal sealed class ShrinkSchedulerQueue : IDisposable
    {
        private readonly ConcurrentQueue<ShrinkSchedulerWorkItem> _queue = new();
        private readonly SemaphoreSlim _slots;
        private readonly ShrinkQueueOverflowPolicy _overflowPolicy;
        private readonly string _name;
        private int _disposed;

        public ShrinkSchedulerQueue(string name, ShrinkBusOptions options)
        {
            _name = name;
            _overflowPolicy = options.OverflowPolicy;
            _slots = new SemaphoreSlim(options.QueueCapacity, options.QueueCapacity);
        }

        public bool IsEmpty => _queue.IsEmpty;

        public bool TryEnqueue(ShrinkSchedulerWorkItem item)
        {
            if (Volatile.Read(ref _disposed) != 0)
                return false;

            if (!_slots.Wait(0))
            {
                if (_overflowPolicy != ShrinkQueueOverflowPolicy.DropOldest ||
                    !_queue.TryDequeue(out var dropped))
                    return false;

                dropped.Drop(_name);
                _slots.Release();
                if (!_slots.Wait(0))
                    return false;
            }

            _queue.Enqueue(item);
            return true;
        }

        public async UniTask EnqueueAsync(ShrinkSchedulerWorkItem item, CancellationToken cancellationToken)
        {
            if (Volatile.Read(ref _disposed) != 0)
                throw new ObjectDisposedException(_name);

            if (_overflowPolicy == ShrinkQueueOverflowPolicy.Wait)
            {
                await _slots.WaitAsync(cancellationToken).AsUniTask(useCurrentSynchronizationContext: false);
                if (Volatile.Read(ref _disposed) != 0)
                {
                    _slots.Release();
                    throw new ObjectDisposedException(_name);
                }
                _queue.Enqueue(item);
                return;
            }

            if (!TryEnqueue(item))
                throw new ShrinkQueueFullException(_name);
        }

        public bool TryDequeue(out ShrinkSchedulerWorkItem item)
        {
            if (!_queue.TryDequeue(out item!))
                return false;
            _slots.Release();
            return true;
        }

        public void DropPending()
        {
            while (_queue.TryDequeue(out var item))
            {
                _slots.Release();
                item.Drop(_name);
            }
        }

        public void Dispose()
        {
            if (Interlocked.Exchange(ref _disposed, 1) != 0)
                return;
            DropPending();
            _slots.Dispose();
        }
    }

    internal static class ShrinkSchedulerWorkItemRunner
    {
        public static async UniTask RunAsync(ShrinkSchedulerWorkItem item)
        {
            try
            {
                await item.Action();
                item.Complete();
            }
            catch (OperationCanceledException ex)
            {
                item.Cancel(ex.CancellationToken);
            }
            catch (Exception ex)
            {
                item.Fail(ex);
                ShrinkEventDiagnostics.LogException(ex);
            }
        }

        public static void RunBlocking(ShrinkSchedulerWorkItem item)
        {
            try
            {
                item.Action().GetAwaiter().GetResult();
                item.Complete();
            }
            catch (OperationCanceledException ex)
            {
                item.Cancel(ex.CancellationToken);
            }
            catch (Exception ex)
            {
                item.Fail(ex);
                ShrinkEventDiagnostics.LogException(ex);
            }
        }
    }

    internal sealed class ShrinkInlineScheduler : IShrinkBusScheduler
    {
        public bool IsOnSchedulerThread => true;

        public bool TryPost(Action action)
        {
            action();
            return true;
        }

        public UniTask PostAsync(Func<UniTask> action, CancellationToken cancellationToken)
        {
            cancellationToken.ThrowIfCancellationRequested();
            return action();
        }

        public UniTask ShutdownAsync(bool drain, TimeSpan timeout) => UniTask.CompletedTask;
        public void Dispose() { }
    }

    internal sealed class ShrinkMainThreadScheduler : IShrinkBusScheduler
    {
        private readonly ShrinkSchedulerQueue _queue;
        private int _pumpScheduled;
        private int _disposed;

        public ShrinkMainThreadScheduler(string name, ShrinkBusOptions options)
        {
            _queue = new ShrinkSchedulerQueue(name, options);
        }

        public bool IsOnSchedulerThread => PlayerLoopHelper.IsMainThread;

        public bool TryPost(Action action)
        {
            if (Volatile.Read(ref _disposed) != 0)
                return false;
            if (IsOnSchedulerThread)
            {
                action();
                return true;
            }

            var item = new ShrinkSchedulerWorkItem(() =>
            {
                action();
                return UniTask.CompletedTask;
            }, awaitable: false);
            if (!_queue.TryEnqueue(item))
                return false;
            SchedulePump();
            return true;
        }

        public async UniTask PostAsync(Func<UniTask> action, CancellationToken cancellationToken)
        {
            if (Volatile.Read(ref _disposed) != 0)
                throw new ObjectDisposedException(nameof(ShrinkMainThreadScheduler));
            cancellationToken.ThrowIfCancellationRequested();
            if (IsOnSchedulerThread)
            {
                await action();
                return;
            }

            var item = new ShrinkSchedulerWorkItem(action, awaitable: true);
            await _queue.EnqueueAsync(item, cancellationToken);
            SchedulePump();
            await item.Completion.AttachExternalCancellation(cancellationToken);
        }

        public async UniTask ShutdownAsync(bool drain, TimeSpan timeout)
        {
            Interlocked.Exchange(ref _disposed, 1);
            if (!drain)
                _queue.DropPending();
            else
                await ShrinkSchedulerShutdown.WaitUntilDrainedAsync(
                    _queue, () => Volatile.Read(ref _pumpScheduled) != 0, timeout);
            _queue.Dispose();
        }

        public void Dispose()
        {
            if (Interlocked.Exchange(ref _disposed, 1) != 0)
                return;
            _queue.Dispose();
        }

        private void SchedulePump()
        {
            if (Interlocked.Exchange(ref _pumpScheduled, 1) != 0)
                return;
            UniTask.Void(PumpAsync);
        }

        private async UniTaskVoid PumpAsync()
        {
            await UniTask.SwitchToMainThread();
            try
            {
                while (_queue.TryDequeue(out var item))
                    await ShrinkSchedulerWorkItemRunner.RunAsync(item);
            }
            finally
            {
                Interlocked.Exchange(ref _pumpScheduled, 0);
                if (!_queue.IsEmpty && Volatile.Read(ref _disposed) == 0)
                    SchedulePump();
            }
        }
    }

    internal sealed class ShrinkTaskPoolScheduler : IShrinkBusScheduler
    {
        private readonly ShrinkSchedulerQueue _queue;
        private readonly int _maxConcurrency;
        private int _workers;
        private int _disposed;

        public ShrinkTaskPoolScheduler(string name, ShrinkBusOptions options)
        {
            _queue = new ShrinkSchedulerQueue(name, options);
            _maxConcurrency = options.DispatchMode == ShrinkDispatchMode.Ordered ? 1 : options.MaxConcurrency;
        }

        public bool IsOnSchedulerThread => false;

        public bool TryPost(Action action)
        {
            if (Volatile.Read(ref _disposed) != 0)
                return false;
            var item = new ShrinkSchedulerWorkItem(() =>
            {
                action();
                return UniTask.CompletedTask;
            }, awaitable: false);
            if (!_queue.TryEnqueue(item))
                return false;
            ScheduleWorkers();
            return true;
        }

        public async UniTask PostAsync(Func<UniTask> action, CancellationToken cancellationToken)
        {
            if (Volatile.Read(ref _disposed) != 0)
                throw new ObjectDisposedException(nameof(ShrinkTaskPoolScheduler));
            var item = new ShrinkSchedulerWorkItem(action, awaitable: true);
            await _queue.EnqueueAsync(item, cancellationToken);
            ScheduleWorkers();
            await item.Completion.AttachExternalCancellation(cancellationToken);
        }

        public async UniTask ShutdownAsync(bool drain, TimeSpan timeout)
        {
            Interlocked.Exchange(ref _disposed, 1);
            if (!drain)
                _queue.DropPending();
            else
                await ShrinkSchedulerShutdown.WaitUntilDrainedAsync(
                    _queue, () => Volatile.Read(ref _workers) != 0, timeout);
            _queue.Dispose();
        }

        public void Dispose()
        {
            if (Interlocked.Exchange(ref _disposed, 1) != 0)
                return;
            _queue.Dispose();
        }

        private void ScheduleWorkers()
        {
            while (!_queue.IsEmpty)
            {
                var current = Volatile.Read(ref _workers);
                if (current >= _maxConcurrency ||
                    Interlocked.CompareExchange(ref _workers, current + 1, current) != current)
                    return;
                UniTask.Void(WorkerAsync);
            }
        }

        private async UniTaskVoid WorkerAsync()
        {
            await UniTask.SwitchToThreadPool();
            try
            {
                while (_queue.TryDequeue(out var item))
                    await ShrinkSchedulerWorkItemRunner.RunAsync(item);
            }
            finally
            {
                Interlocked.Decrement(ref _workers);
                if (!_queue.IsEmpty && Volatile.Read(ref _disposed) == 0)
                    ScheduleWorkers();
            }
        }
    }

    internal sealed class ShrinkDedicatedThreadScheduler : IShrinkBusScheduler
    {
        private readonly ShrinkSchedulerQueue _queue;
        private readonly AutoResetEvent _signal = new(false);
        private readonly ManualResetEventSlim _stopped = new(false);
        private readonly Thread _thread;
        private int _disposed;
        private int _threadId;
        private int _shutdownTimedOut;
        private int _waitHandlesDisposed;

        public ShrinkDedicatedThreadScheduler(string name, ShrinkBusOptions options)
        {
            _queue = new ShrinkSchedulerQueue(name, options);
            _thread = new Thread(Run)
            {
                IsBackground = true,
                Name = $"ShrinkBus:{name}"
            };
            _thread.Start();
        }

        public bool IsOnSchedulerThread => Thread.CurrentThread.ManagedThreadId == Volatile.Read(ref _threadId);

        public bool TryPost(Action action)
        {
            if (Volatile.Read(ref _disposed) != 0)
                return false;
            if (IsOnSchedulerThread)
            {
                action();
                return true;
            }

            var item = new ShrinkSchedulerWorkItem(() =>
            {
                action();
                return UniTask.CompletedTask;
            }, awaitable: false);
            if (!_queue.TryEnqueue(item))
                return false;
            _signal.Set();
            return true;
        }

        public async UniTask PostAsync(Func<UniTask> action, CancellationToken cancellationToken)
        {
            if (Volatile.Read(ref _disposed) != 0)
                throw new ObjectDisposedException(nameof(ShrinkDedicatedThreadScheduler));
            if (IsOnSchedulerThread)
            {
                await action();
                return;
            }

            var item = new ShrinkSchedulerWorkItem(action, awaitable: true);
            await _queue.EnqueueAsync(item, cancellationToken);
            _signal.Set();
            await item.Completion.AttachExternalCancellation(cancellationToken);
        }

        public async UniTask ShutdownAsync(bool drain, TimeSpan timeout)
        {
            Interlocked.Exchange(ref _disposed, 1);
            if (!drain)
                _queue.DropPending();
            _signal.Set();
            var waitMs = timeout == Timeout.InfiniteTimeSpan
                ? Timeout.Infinite
                : Math.Max(0, (int)Math.Min(int.MaxValue, timeout.TotalMilliseconds));
            var stopped = await UniTask.RunOnThreadPool(() => _stopped.Wait(waitMs));
            if (!stopped)
            {
                _queue.DropPending();
                Volatile.Write(ref _shutdownTimedOut, 1);
                try
                {
                    if (_stopped.IsSet)
                        DisposeWaitHandles();
                }
                catch (ObjectDisposedException)
                {
                }
                return;
            }
            DisposeWaitHandles();
        }

        public void Dispose()
        {
            if (Interlocked.Exchange(ref _disposed, 1) != 0)
                return;
            _queue.DropPending();
            Volatile.Write(ref _shutdownTimedOut, 1);
            _signal.Set();
        }

        private void Run()
        {
            Volatile.Write(ref _threadId, Thread.CurrentThread.ManagedThreadId);
            try
            {
                while (Volatile.Read(ref _disposed) == 0 || !_queue.IsEmpty)
                {
                    if (!_queue.TryDequeue(out var item))
                    {
                        _signal.WaitOne(50);
                        continue;
                    }
                    ShrinkSchedulerWorkItemRunner.RunBlocking(item);
                }
            }
            finally
            {
                _stopped.Set();
                if (Volatile.Read(ref _shutdownTimedOut) != 0)
                    DisposeWaitHandles();
            }
        }

        private void DisposeWaitHandles()
        {
            if (Interlocked.Exchange(ref _waitHandlesDisposed, 1) != 0)
                return;
            _queue.Dispose();
            _signal.Dispose();
            _stopped.Dispose();
        }
    }

    internal static class ShrinkEventDiagnostics
    {
        public static void LogException(Exception exception)
        {
#if UNITY_5_3_OR_NEWER
            UnityEngine.Debug.LogException(exception);
#else
            Trace.TraceError(exception.ToString());
#endif
        }
    }

    internal static class ShrinkSchedulerShutdown
    {
        public static async UniTask WaitUntilDrainedAsync(ShrinkSchedulerQueue queue,
            Func<bool> hasActiveWork, TimeSpan timeout)
        {
            var started = Stopwatch.StartNew();
            while (!queue.IsEmpty || hasActiveWork())
            {
                if (timeout != Timeout.InfiniteTimeSpan && started.Elapsed >= timeout)
                {
                    queue.DropPending();
                    return;
                }
                await UniTask.Delay(1, ignoreTimeScale: true);
            }
        }
    }

}