using ClaudeDo.Worker.Runner;
using Xunit;

namespace ClaudeDo.Worker.Tests.Runner;

public class PendingQuestionRegistryTests
{
    [Fact]
    public async Task Register_ThenAnswer_CompletesTheWait()
    {
        var registry = new PendingQuestionRegistry();
        var (questionId, answer) = registry.Register("t1", "which?");

        Assert.False(answer.IsCompleted);
        Assert.Equal("which?", registry.Get("t1")?.Question);

        Assert.True(registry.TryAnswer("t1", questionId, "this one"));
        Assert.Equal("this one", await answer);
        Assert.Null(registry.Get("t1"));            // cleared after answering
    }

    [Fact]
    public void TryAnswer_WrongQuestionId_DoesNothing()
    {
        var registry = new PendingQuestionRegistry();
        var (_, answer) = registry.Register("t1", "q?");

        Assert.False(registry.TryAnswer("t1", "stale-id", "x"));
        Assert.False(answer.IsCompleted);
        Assert.NotNull(registry.Get("t1"));
    }

    [Fact]
    public void TryAnswer_UnknownTask_ReturnsFalse()
    {
        var registry = new PendingQuestionRegistry();
        Assert.False(registry.TryAnswer("ghost", "q", "x"));
    }

    [Fact]
    public void SecondRegister_OverwritesStaleEntry()
    {
        var registry = new PendingQuestionRegistry();
        var (firstId, _) = registry.Register("t1", "first");
        var (secondId, _) = registry.Register("t1", "second");

        Assert.NotEqual(firstId, secondId);
        Assert.Equal("second", registry.Get("t1")?.Question);
        Assert.False(registry.TryAnswer("t1", firstId, "x"));   // old id no longer valid
        Assert.True(registry.TryAnswer("t1", secondId, "ok"));
    }

    [Fact]
    public void Remove_DropsPendingWithoutAnswering()
    {
        var registry = new PendingQuestionRegistry();
        var (questionId, answer) = registry.Register("t1", "q?");

        registry.Remove("t1", questionId);

        Assert.Null(registry.Get("t1"));
        Assert.False(answer.IsCompleted);
    }
}

public class AskUserWaitTests
{
    [Fact]
    public async Task AwaitAnswer_ReturnsUserAnswer_WhenAnsweredInTime()
    {
        var registry = new PendingQuestionRegistry();
        string? askedQuestionId = null;
        var asked = new TaskCompletionSource();
        var resolved = 0;

        var wait = TaskRunMcpService.AwaitAnswerAsync(
            registry, "t1", "DPAPI or plaintext?",
            onAsked: (qid, _) => { askedQuestionId = qid; asked.TrySetResult(); return Task.CompletedTask; },
            onResolved: _ => { resolved++; return Task.CompletedTask; },
            timeout: TimeSpan.FromSeconds(5),
            ct: CancellationToken.None);

        await asked.Task;   // registration + onAsked have run
        Assert.True(registry.TryAnswer("t1", askedQuestionId!, "DPAPI please"));

        Assert.Equal("DPAPI please", await wait);
        Assert.Equal(1, resolved);
        Assert.Null(registry.Get("t1"));
    }

    [Fact]
    public async Task AwaitAnswer_ReturnsFallback_OnTimeout()
    {
        var registry = new PendingQuestionRegistry();
        var resolved = 0;

        var result = await TaskRunMcpService.AwaitAnswerAsync(
            registry, "t2", "q?",
            onAsked: (_, _) => Task.CompletedTask,
            onResolved: _ => { resolved++; return Task.CompletedTask; },
            timeout: TimeSpan.FromMilliseconds(40),
            ct: CancellationToken.None);

        Assert.Equal(TaskRunMcpService.TimeoutFallback, result);
        Assert.Equal(1, resolved);
        Assert.Null(registry.Get("t2"));        // cleaned up after timeout
    }

    [Fact]
    public async Task AwaitAnswer_Rethrows_WhenRunCancelled()
    {
        var registry = new PendingQuestionRegistry();
        using var cts = new CancellationTokenSource();
        cts.Cancel();

        await Assert.ThrowsAnyAsync<OperationCanceledException>(() =>
            TaskRunMcpService.AwaitAnswerAsync(
                registry, "t3", "q?",
                onAsked: (_, _) => Task.CompletedTask,
                onResolved: _ => Task.CompletedTask,
                timeout: TimeSpan.FromMinutes(1),
                ct: cts.Token));

        Assert.Null(registry.Get("t3"));        // cleanup still ran
    }
}