using System.Collections.Concurrent;

namespace ClaudeDo.Worker.Queue;

/// Maps a running task id to the CancellationTokenSource driving its Claude process.
/// QueueService and OverrideSlotService register their slots here so components that
/// cancel tasks DB-side (TaskStateService.CancelAsync and its child cascade) can also
/// stop the process — without depending on the queue services, which would create a
/// DI cycle (QueueService → TaskRunner → ITaskStateService).
public sealed class RunCancellationRegistry
{
    private readonly ConcurrentDictionary<string, CancellationTokenSource> _running = new(StringComparer.Ordinal);

    public void Register(string taskId, CancellationTokenSource cts) => _running[taskId] = cts;

    /// Removes the registration only if <paramref name="cts"/> is still the registered
    /// one — a re-run may already have registered a newer CTS under the same task id.
    public void Unregister(string taskId, CancellationTokenSource cts) =>
        _running.TryRemove(new KeyValuePair<string, CancellationTokenSource>(taskId, cts));

    public bool TryCancel(string taskId)
    {
        if (!_running.TryGetValue(taskId, out var cts)) return false;
        try
        {
            cts.Cancel();
            return true;
        }
        catch (ObjectDisposedException)
        {
            // Slot cleanup raced us; the run is already finished.
            return false;
        }
    }
}