#!/usr/bin/env bash
# description: Spawn N subagents with a shared prompt, wait for all to complete, return concatenated results (fork-join).
# Args: [-n N] [-agent A] [-backend B] [-model M] [-cwd DIR] [-keep] <prompt...>
# Output: concatenated results from all subagents, in submission order. With -n > 1, results are separated by === {id} === headers.
set -euo pipefail

OLLIE="${OLLIE:-$HOME/mnt/ollie}"

die() { printf 'subagent_spawn: %s\n' "$*" >&2; exit 1; }

n=1; agent=; backend=; model=; cwd=; keep=0

while [ $# -gt 0 ]; do
    case $1 in
        -n)       n=$2;       shift ;;
        -agent)   agent=$2;   shift ;;
        -backend) backend=$2; shift ;;
        -model)   model=$2;   shift ;;
        -cwd)     cwd=$2;     shift ;;
        -keep)    keep=1 ;;
        --)       shift; break ;;
        -*)       die "unknown flag: $1" ;;
        *)        break ;;
    esac
    shift
done

[ $# -eq 0 ] && die "prompt required"
prompt="$*"
cwd="${cwd:-$(pwd)}"

base="spawn-$(date +%s%N)-$$"

# Submit all jobs as one parallel spec.
{
    printf 'name=%s\n'     "$base"
    printf 'cwd=%s\n'      "$cwd"
    [ -n "$agent"   ] && printf 'agent=%s\n'   "$agent"
    [ -n "$backend" ] && printf 'backend=%s\n' "$backend"
    [ -n "$model"   ] && printf 'model=%s\n'   "$model"
    printf 'parallel=%s\n' "$n"
    printf -- '---\n'
    printf '%s\n' "$prompt"
} > "$OLLIE/b/new" || die "failed to submit jobs"

# Fork: poll each job's status in a background subshell, write result to tmpdir.
tmpdir=$(mktemp -d)
trap 'rm -rf "$tmpdir"' EXIT

declare -a pids
for i in $(seq 0 $((n - 1))); do
    id="${base}-${i}"
    (
        while true; do
            status=$(tr -d '[:space:]' < "$OLLIE/b/$id/status" 2>/dev/null) || { sleep 0.05; continue; }
            case "$status" in
                done*)   cat "$OLLIE/b/$id/result" 2>/dev/null > "$tmpdir/$i"; break ;;
                failed*) printf 'error: subagent %s failed: %s\n' "$id" "$status" > "$tmpdir/$i"; break ;;
                *)       sleep 0.05 ;;
            esac
        done
    ) &
    pids[$i]=$!
done

# Join: wait for all pollers.
for i in $(seq 0 $((n - 1))); do
    wait "${pids[$i]}"
done

# Collect results in submission order.
for i in $(seq 0 $((n - 1))); do
    id="${base}-${i}"
    [ "$n" -gt 1 ] && printf '=== %s ===\n' "$id"
    cat "$tmpdir/$i" 2>/dev/null || true
    [ "$keep" -eq 0 ] && rm -r "$OLLIE/b/$id" 2>/dev/null || true
done