#!/usr/bin/env bash
# s/b: one-shot ollie query via s/ namespace.
#
# Usage: job [flags] [prompt...]
#        echo "prompt" | job [flags]

set -euo pipefail

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

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

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

Flags:
  -name NAME        session name
  -cwd DIR          working directory (default: pwd)
  -agent AGENT      agent config name
  -backend BACKEND  backend override (e.g. ollama, anthropic)
  -model MODEL      model override
  -json             append JSON-only instruction to prompt
  -parallel N       run N times in parallel (default: 1)
  -keep             do not kill session(s) after printing result
  -bg               submit and exit immediately; prints session path(s); implies -keep
  -failfast         on first failure, cancel remaining parallel jobs and exit
  -help             show this help
EOF
    exit 0
}

name=; cwd=; agent=; backend=; model=; json=0; 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)      json=1 ;;
        -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

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"

if [ "$json" -eq 1 ]; then
    prompt="$prompt"$'\n\n'"Respond with valid JSON only, no prose."
fi

cwd=${cwd:-$(pwd)}
base=${name:-$(date +%s%N)-$$}

fn_run_one() {
    local id=$1
    local B="$MOUNT/s/$id"

    # Build as a single string: bash's printf builtin issues one write() per call;
    # the server processes each write independently, so partial KV would fail.
    local kv="name=${id}
cwd=${cwd}"
    [ -n "$agent"   ] && kv="${kv}
agent=${agent}"
    [ -n "$backend" ] && kv="${kv}
backend=${backend}"
    [ -n "$model"   ] && kv="${kv}
model=${model}"
    printf '%s\n' "$kv" > "$MOUNT/s/new" || die "failed to create session"

    # Poll until the session directory is visible (FUSE entry cache may lag).
    local deadline=$(( $(date +%s) + 5 ))
    until [ -d "$B" ] || [ "$(date +%s)" -ge "$deadline" ]; do sleep 0.05; done
    [ -d "$B" ] || die "timeout waiting for session $id"

    printf '%s\n' "$prompt" > "$B/prompt" || die "failed to submit prompt"

    if [ "$bg" -eq 1 ]; then
        printf '%s\n' "$B"
        return 0
    fi

    # Wait until idle.
    while true; do
        state=$(cat "$B/statewait" 2>/dev/null | tr -d '[:space:]') || break
        [ "$state" = "idle" ] && break
    done

    local rc=0
    case "$state" in
        failed*|error*) rc=1 ;;
    esac

    # Print result: chat content from offset onward.
    local offset
    offset=$(cat "$B/offset" 2>/dev/null || echo 0)
    tail -c +$((offset + 1)) "$B/chat" 2>/dev/null
    printf '\n'

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

rc=0

if [ "$parallel" -gt 1 ]; then
    abort=$(mktemp)
    trap 'rm -f "$abort"' EXIT
    pids=()
    for i in $(seq 1 "$parallel"); do
        id="${base}-${i}"
        (
            [ -s "$abort" ] && exit 1
            fn_run_one "$id" || {
                if [ "$failfast" -eq 1 ]; then
                    echo 1 > "$abort"
                    for j in $(seq 1 "$parallel"); do
                        jid="${base}-${j}"
                        [ "$jid" != "$id" ] && rm -r "$MOUNT/s/$jid" 2>/dev/null || true
                    done
                fi
                exit 1
            }
        ) &
        pids+=($!)
    done
    for pid in "${pids[@]}"; do wait "$pid" || rc=1; done
else
    fn_run_one "$base" || rc=1
fi

exit $rc