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);
    private readonly ILogger<RunCancellationRegistry> _logger;

    public RunCancellationRegistry(ILogger<RunCancellationRegistry> logger) => _logger = logger;

    /// Registers the CTS driving <paramref name="taskId"/>'s run. Refuses (and logs) instead of
    /// silently overwriting when a CTS is already registered for this task id — an overwrite
    /// would mean a double-dispatch is in flight, and the caller that lost the race would then
    /// unregister the winner's CTS out from under it, leaving CancelAsync unable to reach the
    /// still-running process.
    public bool Register(string taskId, CancellationTokenSource cts)
    {
        if (_running.TryAdd(taskId, cts)) return true;

        _logger.LogWarning(
            "Task {TaskId} already has a registered run; refusing to overwrite (double-dispatch?)", taskId);
        return false;
    }

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