#!/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; }

# Retry a command once after remounting on I/O error.
try() { "$@" 2>/dev/null || { ollie-remount && "$@"; }; }
# Write $1 to file $2, retrying after remount on failure.
wf() { printf '%s\n' "$1" > "$2" 2>/dev/null || { ollie-remount && printf '%s\n' "$1" > "$2"; }; }

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)-$$}

# JSON-escape a string for embedding in a JSON value.
jsonesc() {
    printf '%s' "$1" | sed 's/\\/\\\\/g; s/"/\\"/g; s/\t/\\t/g' | tr '\n' '\1' | sed 's/\x01/\\n/g'
}

# Foreground generation via /generate (net/dns pattern).
fn_generate() {
    local json='{"prompt":"'$(jsonesc "$prompt")'"'
    [ -n "$backend" ] && json="$json,\"backend\":\"$backend\""
    [ -n "$model" ]   && json="$json,\"model\":\"$model\""
    json="$json}"

    printf '%s' "$json" | 9p rdwr ollie/generate || die "generate failed (is olliesrv running?)"
}

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}"
    wf "$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"

    wf "$prompt" "$B/prompt" || die "failed to submit prompt"

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

    # Poll state until idle (avoids blocking FUSE reads that can deadlock
    # when parallel sessions access the same mount).
    local state
    while true; do
        state=$(try cat "$B/state" | tr -d '[:space:]') || break
        [ "$state" = "idle" ] && break
        sleep 0.5
    done

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

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

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

rc=0

if [ "$bg" -eq 1 ]; then
    # Background mode: must use sessions (persistent state to poll later).
    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
else
    # Foreground mode: use /generate (stateless, no session).
    if [ "$parallel" -gt 1 ]; then
        pids=()
        for i in $(seq 1 "$parallel"); do
            fn_generate &
            pids+=($!)
        done
        for pid in "${pids[@]}"; do wait "$pid" || rc=1; done
    else
        fn_generate || rc=1
    fi
fi

exit $rc