using Microsoft.Agents.AI.Workflows;

// Repro: a fan-in barrier (AddFanInBarrierEdge) loses a contribution that was
// buffered BEFORE a checkpoint/resume, so after resume the barrier never
// completes and its target never runs.
//
//   A ──contributes "A"──────────────────────────┐ (buffered before the pause)
//   └──asks the human (RequestPort) ──suspend──> B ┤ FAN-IN BARRIER ──> Sink
//                                  (reply)        (contributes "B" after resume)
//
// Expected: after answering the request, the barrier releases both A and B to Sink.
// Actual:   only B is produced after resume; A was dropped, the barrier never
//           completes, and Sink never runs.

static Workflow Build()
{
    var a = new A();
    var b = new B();
    var sink = new Sink();
    var port = RequestPort.Create<Ask, Reply>("approval");

    var builder = new WorkflowBuilder(a);
    builder.AddEdge(a, port);     // A asks the human -> run suspends
    builder.AddEdge(port, b);     // human reply -> B
    builder.AddFanInBarrierEdge(new ExecutorBinding[] { a, b }, sink); // barrier: waits for A and B
    return builder.Build();
}

var manager = CheckpointManager.CreateInMemory();
var env = InProcessExecution.Lockstep.WithCheckpointing(manager);

// 1) Run until the human-input request suspends the workflow; grab that checkpoint.
CheckpointInfo? atSuspend = null;
bool sawRequest = false;
await using (var run = await env.RunStreamingAsync(Build(), "start", "session-1"))
{
    await foreach (var evt in run.WatchStreamAsync())
    {
        if (evt is RequestInfoEvent) sawRequest = true;
        if (evt is SuperStepCompletedEvent { CompletionInfo.Checkpoint: { } cp }) atSuspend = cp;
        if (sawRequest && atSuspend is not null) break;   // suspended with a checkpoint in hand
    }
}
if (atSuspend is null) { Console.WriteLine("inconclusive: never suspended"); return 2; }

// 2) Resume from that checkpoint and answer the request so B contributes.
Sink.Received.Clear();
await using (var run = await env.ResumeStreamingAsync(Build(), atSuspend))
{
    var answered = false;
    await foreach (var evt in run.WatchStreamAsync())
    {
        if (!answered && evt is RequestInfoEvent req)
        {
            await run.SendResponseAsync(req.Request.CreateResponse(new Reply("yes")));
            answered = true;
        }
        else if (answered && evt is SuperStepCompletedEvent { CompletionInfo: { HasPendingMessages: false, HasPendingRequests: false } })
            break;
    }
}

// 3) Did the barrier release both contributions to Sink?
Console.WriteLine($"Sink received after resume: {{{string.Join(",", Sink.Received.OrderBy(x => x))}}}");
var ok = Sink.Received.SetEquals(new[] { "A", "B" });
Console.WriteLine(ok
    ? "PASS: barrier completed (A and B both reached Sink)."
    : "BUG: A was buffered before the pause but lost across resume; the barrier never completed and Sink never ran.");
return ok ? 0 : 1;

record Ask(string Prompt);
record Reply(string Answer);
record Item(string Source);

// A is the start executor AND the first barrier source: it contributes immediately,
// then asks the human in the same step (which suspends the run).
sealed class A() : Executor("A")
{
    protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder pb) =>
        pb.ConfigureRoutes(r => r.AddHandler<string>(OnStart)).SendsMessage<Item>().SendsMessage<Ask>();

    private async ValueTask OnStart(string input, IWorkflowContext ctx)
    {
        Console.WriteLine("[A] contributing 'A' to the barrier, then asking the human (suspends)");
        await ctx.SendMessageAsync(new Item("A"));          // -> fan-in barrier (buffered)
        await ctx.SendMessageAsync(new Ask("approve?"));    // -> RequestPort (suspends)
    }
}

// B is the second barrier source: it contributes only after the human replies.
sealed class B() : Executor("B")
{
    protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder pb) =>
        pb.ConfigureRoutes(r => r.AddHandler<Reply>(OnReply)).SendsMessage<Item>();

    private ValueTask OnReply(Reply reply, IWorkflowContext ctx)
    {
        Console.WriteLine("[B] contributing 'B' to the barrier (after resume)");
        return ctx.SendMessageAsync(new Item("B"));
    }
}

// Sink is the fan-in barrier target; the framework runs it once per contribution when the barrier completes.
sealed class Sink() : Executor("sink")
{
    public static readonly HashSet<string> Received = new();

    protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder pb) =>
        pb.ConfigureRoutes(r => r.AddHandler<Item>(OnItem)).YieldsOutput<string>();

    private ValueTask OnItem(Item item, IWorkflowContext ctx)
    {
        Received.Add(item.Source);
        return default;
    }
}