#!/usr/bin/env bash
# u/optimize: generate N candidate optimized prompts and select the best.
set -euo pipefail

for _arg; do
    [ "$_arg" = "-help" ] && cat <<'EOF' && exit 0
usage: optimize [flags] [prompt...]
       echo "prompt" | optimize [flags]

Generates N optimized candidates for a prompt, then judges and returns the best.
Uses OLLIE_MODEL and OLLIE_BACKEND by default for both stages.

Flags:
  -n N          number of candidates to generate (default: 3)
  -cb BACKEND   candidate generation backend
  -cm MODEL     candidate generation model
  -jb BACKEND   judge backend
  -jm MODEL     judge model
  -help         show this help
EOF
done

n=3; cb=; cm=; jb=; jm=

while [ $# -gt 0 ]; do
    case $1 in
        -n)    n=$2;  shift ;;
        -cb)   cb=$2; shift ;;
        -cm)   cm=$2; shift ;;
        -jb)   jb=$2; shift ;;
        -jm)   jm=$2; shift ;;
        -help) ;;
        --)    shift; break ;;
        -*)    printf 'optimize: unknown flag: %s\n' "$1" >&2; exit 1 ;;
        *)     break ;;
    esac
    shift
done

if [ $# -gt 0 ]; then
    prompt="$*"
elif ! [ -t 0 ]; then
    prompt=$(cat)
else
    printf 'optimize: prompt required\n' >&2; exit 1
fi

[ -z "$prompt" ] && { printf 'optimize: empty prompt\n' >&2; exit 1; }

Q="${0%/*}/../b/q"

c_flags=()
j_flags=()
[ -n "$cb" ] && c_flags+=(-backend "$cb")
[ -n "$cm" ] && c_flags+=(-model "$cm")
[ -n "$jb" ] && j_flags+=(-backend "$jb")
[ -n "$jm" ] && j_flags+=(-model "$jm")

printf 'generating %d candidates...\n' "$n" >&2
candidates=$("$Q" "${c_flags[@]}" -parallel "$n" \
    "Optimize the following prompt. Return only the optimized prompt, no explanation or preamble:

$prompt")

printf 'judging...\n' >&2
"$Q" "${j_flags[@]}" \
    "Review these candidate prompts and return only the single best one, no explanation or preamble:

$candidates"