using System.Collections.Concurrent;

namespace ClaudeDo.Worker.Runner;

public sealed record PendingQuestion(string TaskId, string QuestionId, string Question);

// In-memory store of questions a running task has raised via the AskUser MCP tool and is
// blocking on. One pending question per task (the run's process is blocked mid-tool-call,
// so it cannot ask twice at once). Kept out of the DB on purpose: a question that outlives
// a Worker restart is already dead (StaleTaskRecovery flips the run to Failed). Singleton.
public sealed class PendingQuestionRegistry
{
    private readonly ConcurrentDictionary<string, Entry> _byTask = new();

    private sealed record Entry(string QuestionId, string Question, TaskCompletionSource<string> Answer);

    // Registers a question for the task and returns its id plus the awaitable answer.
    // A second register for the same task replaces any stale entry.
    public (string QuestionId, Task<string> Answer) Register(string taskId, string question)
    {
        var questionId = Guid.NewGuid().ToString("N");
        var tcs = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
        _byTask[taskId] = new Entry(questionId, question, tcs);
        return (questionId, tcs.Task);
    }

    // Delivers the answer to a waiting question. Returns false if no matching question is
    // pending (already answered, timed out, or stale id).
    public bool TryAnswer(string taskId, string questionId, string answer)
    {
        if (_byTask.TryGetValue(taskId, out var entry)
            && entry.QuestionId == questionId
            && _byTask.TryRemove(taskId, out _))
        {
            return entry.Answer.TrySetResult(answer ?? string.Empty);
        }
        return false;
    }

    public PendingQuestion? Get(string taskId) =>
        _byTask.TryGetValue(taskId, out var entry)
            ? new PendingQuestion(taskId, entry.QuestionId, entry.Question)
            : null;

    // Drops a pending question without delivering an answer (timeout/cancel cleanup).
    public void Remove(string taskId, string questionId)
    {
        if (_byTask.TryGetValue(taskId, out var entry) && entry.QuestionId == questionId)
            _byTask.TryRemove(taskId, out _);
    }
}