w4c-workflows-api/Services/Nodes/Executors/SplitInBatchesNodeExecutor.cs

using System.Globalization;
using System.Text.Json.Nodes;
using w4c_workflows.Models.Nodes;

namespace w4c_workflows.Services.Nodes.Executors;

/// <summary>
/// Splitting / loop node. Two behaviours share one node:
///
/// <list type="bullet">
/// <item><b>Chunking</b> (no loop-back edge): runs once with all items and emits
/// every batch on output 0 plus the original items on output 1 (<c>done</c>).</item>
/// <item><b>Loop</b> (a loop-back edge targets this node, see
/// <see cref="NodeBlueprint.LoopBack"/>): emits one batch per invocation on
/// output 0 and, once the input is exhausted, emits the original items on
/// output 1 and signals <see cref="NodeExecutionOutcome.LoopComplete"/> so the
/// runner stops iterating. The cursor lives in the run-scoped state bag, so the
/// node itself stays stateless.</item>
/// </list>
/// </summary>
public sealed class SplitInBatchesNodeExecutor : INodeExecutor
{
    public const string DefaultOutputField = "items";

    public string Type => "core.splitInBatches";

    public Task<NodeExecutionOutcome> RunAsync(NodeExecutionContext context, CancellationToken ct)
    {
        var batchSize = NodeValueAccess.ReadInt(context.Parameters, "batchSize", 10);
        if (batchSize < 1)
            return Task.FromResult(NodeExecutionOutcome.Failed("batchSize must be at least 1", "invalid_parameter"));

        var field = NodeValueAccess.ReadString(context.Parameters, "outputField");
        if (string.IsNullOrWhiteSpace(field))
            field = DefaultOutputField;

        if (!context.LoopBackInput)
        {
            var items = context.Input(0);
            return Task.FromResult(Chunk(items, batchSize, field));
        }

        return Task.FromResult(Loop(context, batchSize, field));
    }

    /// <summary>Stateless: emit every batch at once and pass all items through to <c>done</c>.</summary>
    private static NodeExecutionOutcome Chunk(IReadOnlyList<FlowItem> items, int batchSize, string field)
    {
        var batches = new List<FlowItem>();
        for (var offset = 0; offset < items.Count; offset += batchSize)
        {
            var batch = Slice(items, offset, batchSize);
            batches.Add(new FlowItem
            {
                Json = new JsonObject { [field] = BatchJson(batch) },
                Origin = batch[0].Origin,
            });
        }

        return new NodeExecutionOutcome
        {
            Outputs = new List<IReadOnlyList<FlowItem>> { batches, items.ToList() },
        };
    }

    /// <summary>Loop mode: one batch per invocation, then a final <c>done</c>.</summary>
    private static NodeExecutionOutcome Loop(NodeExecutionContext context, int batchSize, string field)
    {
        var key = "loop:" + context.TaskId;
        if (context.State.TryGetValue(key, out var existing) && existing is LoopCursor cursor)
            return cursor.Next(field);

        var seeded = new LoopCursor(context.Input(0).ToList(), batchSize);
        context.State[key] = seeded;
        return seeded.Next(field);
    }

    private static JsonArray BatchJson(IReadOnlyList<FlowItem> batch)
    {
        var array = new JsonArray();
        foreach (var item in batch)
            array.Add(item.Json.DeepClone());
        return array;
    }

    /// <summary>
    /// Copies <paramref name="count"/> items from <paramref name="offset"/>.
    /// Direct indexing, not <c>Skip().Take()</c>, so chunking is O(n) not O(n²).
    /// </summary>
    private static List<FlowItem> Slice(IReadOnlyList<FlowItem> items, int offset, int count)
    {
        var end = Math.Min(offset + count, items.Count);
        var slice = new List<FlowItem>(Math.Max(0, end - offset));
        for (var i = offset; i < end; i++)
            slice.Add(items[i]);
        return slice;
    }

    /// <summary>Per-run iteration cursor for a loop node.</summary>
    private sealed class LoopCursor
    {
        private readonly int _batchSize;

        public LoopCursor(List<FlowItem> items, int batchSize)
        {
            Items = items;
            _batchSize = batchSize;
        }

        public List<FlowItem> Items { get; }

        public int Offset { get; private set; }

        public NodeExecutionOutcome Next(string field)
        {
            if (Offset >= Items.Count)
            {
                return new NodeExecutionOutcome
                {
                    Outputs = new List<IReadOnlyList<FlowItem>>
                    {
                        Array.Empty<FlowItem>(),
                        Items.ToList(),
                    },
                    LoopComplete = true,
                };
            }

            var batch = Slice(Items, Offset, _batchSize);
            Offset += _batchSize;

            var batchItem = new FlowItem
            {
                Json = new JsonObject { [field] = BatchJson(batch) },
                Origin = batch[0].Origin,
            };

            return new NodeExecutionOutcome
            {
                Outputs = new List<IReadOnlyList<FlowItem>>
                {
                    new List<FlowItem> { batchItem },
                    Array.Empty<FlowItem>(),
                },
            };
        }
    }
}