#!/usr/bin/env bash
# oq: 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
#   -output   output format hint (e.g. json)
#   -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

set -euo pipefail

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

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

name=; cwd=; agent=; backend=; model=; output=; parallel=1; keep=0; bg=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 ;;
        -output)   output=$2;   shift ;;
        -parallel) parallel=$2; shift ;;
        -keep)     keep=1 ;;
        -bg)       bg=1; keep=1 ;;
        --)        shift; break ;;
        -*) die "unknown flag: $1" ;;
        *)  break ;;
    esac
    shift
done

# Prompt from remaining args or stdin.
if [ $# -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"
    [ -n "$output"  ] && printf 'output=%s\n'  "$output"
    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
for i in $(seq 0 $((parallel - 1))); do
    id="${name}-${i}"
    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

    [ "$parallel" -gt 1 ] && printf '=== %s ===\n' "$id"
    cat "$B/result" 2>/dev/null

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

exit $rc
