#!/usr/bin/env bash
# b/job: non-interactive ollie query via b/ namespace.
#
# Usage: oq [flags] [prompt...]
#        echo "prompt" | oq [flags]
#
# Flags:
#   -name     base name for the job(s); suffixed with -0, -1, ...
#   -cwd      working directory (default: pwd)
#   -agent    agent config name
#   -backend  backend override
#   -model    model override
#   -json     shorthand for -output json; emit JSONL output
#   -parallel  run the same prompt N times in parallel (default: 1)
#   -keep      do not remove job(s) after printing result
#   -bg        submit job and exit immediately; prints session id(s); implies -keep
#   -failfast  on first failure, cancel remaining running jobs and exit

set -euo pipefail

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

cmd=${0##*/}
die() { printf '%s: %s\n' "$cmd" "$*" >&2; exit 1; }
trap 'printf "\n"' EXIT

usage() {
    cat <<EOF
usage: $cmd [flags] [prompt...]
       echo "prompt" | $cmd [flags]

Flags:
  -name NAME        base job name; jobs suffixed -0, -1, ...
  -cwd DIR          working directory (default: pwd)
  -agent AGENT      agent config name
  -backend BACKEND  backend override (e.g. ollama, anthropic)
  -model MODEL      model override
  -json             emit JSONL: one {"session":...,"result":...} per line
  -parallel N       run the same prompt N times in parallel (default: 1)
  -keep             do not remove job(s) after printing result
  -bg               submit and exit immediately; prints job id(s); implies -keep
  -failfast         on first failure, cancel remaining running jobs and exit
  -help             show this help
EOF
    exit 0
}

name=; cwd=; agent=; backend=; model=; output=; parallel=1; keep=0; bg=0; failfast=0

while [ $# -gt 0 ]; do
    case $1 in
        -name)      name=$2;     shift ;;
        -cwd)       cwd=$2;      shift ;;
        -agent)     agent=$2;    shift ;;
        -backend)   backend=$2;  shift ;;
        -model)     model=$2;    shift ;;
        -json)      output=json ;;
        -parallel)  parallel=$2; shift ;;
        -keep)      keep=1 ;;
        -bg)        bg=1; keep=1 ;;
        -failfast)  failfast=1 ;;
        -help)      usage ;;
        --)         shift; break ;;
        -*) die "unknown flag: $1" ;;
        *)  break ;;
    esac
    shift
done

# Prompt from args, stdin, or both (args as instruction, stdin as content).
if [ $# -gt 0 ] && ! [ -t 0 ]; then
    prompt="$*"$'\n'"$(cat)"
elif [ $# -gt 0 ]; then
    prompt="$*"
elif ! [ -t 0 ]; then
    prompt=$(cat)
else
    die "prompt required: pass as arguments or pipe via stdin"
fi

[ -z "$prompt" ] && die "empty prompt"

cwd=${cwd:-$(pwd)}
[ -z "$name" ] && name=$(date +%s%N)-$$

{
    printf 'name=%s\n'     "$name"
    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' "$parallel"
    printf -- '---\n'
    printf '%s\n' "$prompt"
} > "$MOUNT/b/new" || die "failed to submit job"

if [ "$bg" -eq 1 ]; then
    for i in $(seq 0 $((parallel - 1))); do
        printf '%s-%s\n' "$name" "$i"
    done
    exit 0
fi

rc=0

if [ "$parallel" -gt 1 ]; then
    abort=$(mktemp)
    trap 'rm -f "$abort"' EXIT
    pids=()
    for i in $(seq 0 $((parallel - 1))); do
        id="${name}-${i}"
        B="$MOUNT/b/$id"
        (
            while true; do
                [ -s "$abort" ] && exit 1
                status=$(tr -d '[:space:]' < "$B/status" 2>/dev/null) || { sleep 0.05; continue; }
                case "$status" in
                    done*)   break ;;
                    failed*) printf 'subagent %s failed: %s\n' "$id" "$status" >&2
                             if [ "$failfast" -eq 1 ]; then
                                 echo 1 > "$abort"
                                 for j in $(seq 0 $((parallel - 1))); do
                                     jid="${name}-${j}"
                                     [ "$jid" != "$id" ] && rm -r "$MOUNT/b/$jid" 2>/dev/null || true
                                 done
                             fi
                             [ "$keep" -eq 0 ] && rm -r "$B" 2>/dev/null || true
                             exit 1 ;;
                    *)       sleep 0.05 ;;
                esac
            done
            if [ -s "$abort" ]; then
                [ "$keep" -eq 0 ] && rm -r "$B" 2>/dev/null || true
                exit 1
            fi
            if [ -n "$output" ]; then
                printf '{"session":"%s","result":%s}\n' "$id" \
                    "$(cat "$B/result" 2>/dev/null | jq -Rs .)"
            else
                printf '=== %s ===\nInput: "%s"\nOutput: ' "$id" "$prompt"
                cat "$B/result" 2>/dev/null
                printf '\n'
            fi
            [ "$keep" -eq 0 ] && rm -r "$B" 2>/dev/null || true
        ) &
        pids+=($!)
    done
    for pid in "${pids[@]}"; do wait "$pid" || rc=1; done
else
    id="${name}-0"
    B="$MOUNT/b/$id"

    while true; do
        status=$(tr -d '[:space:]' < "$B/status" 2>/dev/null) || { sleep 0.05; continue; }
        case "$status" in
            done*)   break ;;
            failed*) rc=1; break ;;
            *)       sleep 0.05 ;;
        esac
    done

    if [ -n "$output" ]; then
        printf '{"session":"%s","result":%s}\n' "$id" \
            "$(cat "$B/result" 2>/dev/null | jq -Rs .)"
    else
        cat "$B/result" 2>/dev/null
        printf '\n'
    fi

    [ "$keep" -eq 0 ] && rm -r "$B" 2>/dev/null || true
fi

exit $rc
