using System.Collections.Concurrent;

namespace LootValueEX.Common;

internal class TaskCache
{
    private readonly SemaphoreSlim _semaphore;
    private readonly ConcurrentDictionary<Guid, Task> _taskDict;
    
    internal TaskCache(int maxTasks = 10)
    {
        _semaphore = new SemaphoreSlim(maxTasks);
        _taskDict = new ConcurrentDictionary<Guid, Task>();
    }
    
    private bool IsTaskCancelled(Guid taskGuid) =>
        _taskDict.TryGetValue(taskGuid, out Task task) && task.Status == TaskStatus.Canceled; 

    internal Guid AddTask(Func<CancellationToken, Task> taskFactory)
    {
        Guid taskGuid = Guid.NewGuid();
        CancellationTokenSource cts = new CancellationTokenSource();
        Task task = taskFactory(cts.Token);

        _taskDict[taskGuid] = task;

        return taskGuid;
    }

    internal bool RemoveTask(Guid taskGuid)
    {
        if (_taskDict.TryGetValue(taskGuid, out Task task))
        {
            if (_taskDict.TryRemove(taskGuid, out _))
            {
                task.Dispose();
                _semaphore.Release();
                return true;
            }
        }

        return false;
    }

    internal async Task<Task> RunTaskByGuidAsync(Guid taskGuid)
    {
        await _semaphore.WaitAsync();
        try
        {
            if ((_taskDict.TryGetValue(taskGuid, out Task task)))
            {
                if (!task.IsCompleted && !task.Status.Equals(TaskStatus.Running))
                    await task;
                return task;    
            }
            
            Plugin.Log.LogError($"Task {taskGuid.ToString()} not found.");
        }
        catch (OperationCanceledException)
        {
            RemoveTask(taskGuid);
            Plugin.Log.LogDebug($"Task {taskGuid.ToString()} got cancelled.");
        }
        return Task.FromResult(false);
    }

    internal async Task<Task> AddAndRunTaskAsync(Func<CancellationToken, Task> taskFactory)
    {
        Guid taskGuid = AddTask(taskFactory);
        return await RunTaskByGuidAsync(taskGuid);
    }
}