using System.Collections.Concurrent;
using System.Security.Cryptography;

namespace ClaudeDo.Worker.Runner;

// In-memory per-run MCP identity store. A task run mints a token, registers it here,
// and tears it down when the run ends. Kept out of the DB on purpose: a run that
// outlives a Worker restart is already dead (StaleTaskRecovery flips it to Failed).
public sealed class TaskRunTokenRegistry
{
    private readonly ConcurrentDictionary<string, string> _tokenToTaskId = new();
    public void Register(string token, string taskId) => _tokenToTaskId[token] = taskId;
    public bool TryResolve(string token, out string taskId)
    {
        if (_tokenToTaskId.TryGetValue(token, out var id)) { taskId = id; return true; }
        taskId = string.Empty;
        return false;
    }
    public void Unregister(string token) => _tokenToTaskId.TryRemove(token, out _);

    // Test-only leak check: a run must never leave its token registered after it ends,
    // no matter which step of setup failed.
    public bool HasAnyTokenFor(string taskId) => _tokenToTaskId.Values.Contains(taskId);

    public static string GenerateToken()
    {
        var bytes = RandomNumberGenerator.GetBytes(32);
        return Convert.ToBase64String(bytes).Replace('+', '-').Replace('/', '_').TrimEnd('=');
    }
}
