#!/usr/bin/env bash
# subfleet codex — hardened `codex exec` dispatch.
#
# With no -H, the runner picks a dispatchable lane and moves to another lane
# if the provider reports a usage limit. It also records every dispatch and
# snapshots dirty Git state on exit without changing HEAD, the real index, or
# the worktree. See docs/guard.md for the optional command guard.
#
# Usage:
#   subfleet codex [-H <codex_home>] -m <model> -C <workdir> \
#     -p <prompt_file> -o <out_file> [-s <sandbox>] [-b <salvage_branch>] \
#     [-R <salvage_remote>] [-r <max_retries>] [-e <reasoning_effort>]
set -u

CODEX_HOME_ARG="" MODEL="" WORKDIR="" PROMPT_FILE="" OUT_FILE="" EFFORT=""
SANDBOX="workspace-write" SALVAGE_BRANCH="" SALVAGE_REMOTE="origin" MAX_RETRIES=2
while getopts "H:m:C:p:o:s:b:R:r:e:" opt; do
  case "$opt" in
    H) CODEX_HOME_ARG=$OPTARG ;;
    m) MODEL=$OPTARG ;;
    C) WORKDIR=$OPTARG ;;
    p) PROMPT_FILE=$OPTARG ;;
    o) OUT_FILE=$OPTARG ;;
    s) SANDBOX=$OPTARG ;;
    b) SALVAGE_BRANCH=$OPTARG ;;
    R) SALVAGE_REMOTE=$OPTARG ;;
    r) MAX_RETRIES=$OPTARG ;;
    e) EFFORT=$OPTARG ;;
    *) exit 2 ;;
  esac
done
for req in MODEL WORKDIR PROMPT_FILE OUT_FILE; do
  [ -n "${!req}" ] || { echo "subfleet codex: missing required option" >&2; exit 2; }
done
[ -f "$PROMPT_FILE" ] || {
  echo "subfleet codex: prompt file not found: $PROMPT_FILE" >&2
  exit 2
}
case "$SALVAGE_BRANCH" in
  main|master)
    echo "subfleet codex: refusing -b main|master because salvage pushes are forced; use a dedicated recovery branch" >&2
    exit 2 ;;
esac

ERR_FILE="${OUT_FILE%.md}.err.log"
ORIGINAL_OUT_FILE=${SUBFLEET_RUN_ORIGINAL_OUT-$OUT_FILE}
SCRIPT_DIR=$(CDPATH='' cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
SUBFLEET=${SUBFLEET_RUN_SUBFLEET:-$SCRIPT_DIR/subfleet}
AUTO_LANE=0
RUN_ID=""
EXCLUDED_HOMES=()

record_run_start() {
  local record_args
  record_args=(
    _record-run --phase start --family codex --model "$MODEL"
    --lane "$CODEX_HOME_ARG" --workdir "$WORKDIR" --prompt "$PROMPT_FILE"
    --out "$OUT_FILE" --err "$ERR_FILE" --original-out "$ORIGINAL_OUT_FILE"
  )
  [ -z "${SUBFLEET_RUN_LANE_LOG:-}" ] || record_args+=(--lane-log "$SUBFLEET_RUN_LANE_LOG")
  RUN_ID=$("$SUBFLEET" "${record_args[@]}" 2>/dev/null) || RUN_ID=""
}

record_run_lane() {
  [ -n "$RUN_ID" ] || return 0
  "$SUBFLEET" _record-run --phase update --run-id "$RUN_ID" \
    --lane "$CODEX_HOME_ARG" >/dev/null 2>&1 || true
}

record_run_finish() {
  local final_rc=$1
  local record_args
  [ -n "$RUN_ID" ] || return 0
  record_args=(_record-run --phase finish --run-id "$RUN_ID" --rc "$final_rc")
  [ -z "$CODEX_HOME_ARG" ] || record_args+=(--lane "$CODEX_HOME_ARG")
  "$SUBFLEET" "${record_args[@]}" >/dev/null 2>&1 || true
}

cleanup_owned_prompt() {
  local owned_prompt=${SUBFLEET_RUN_OWNED_PROMPT:-}
  [ -n "$owned_prompt" ] || return 0
  if [ "$owned_prompt" != "$PROMPT_FILE" ]; then
    # A nested invocation may inherit an outer dispatcher marker. It does not
    # own that file, so leave it alone and continue normal cleanup/salvage.
    return 0
  fi
  case "$(basename "$owned_prompt")" in
    delegate-prompt-*.md)
      if ! rm -f "$owned_prompt" || [ -e "$owned_prompt" ]; then
        echo "subfleet codex: WARNING could not remove owned prompt $owned_prompt" >&2
        return 1
      fi ;;
    *)
      echo "subfleet codex: WARNING refusing unexpected owned prompt path $owned_prompt" >&2
      return 1 ;;
  esac
}

early_on_exit() {
  local final_rc=$?
  trap - EXIT
  cleanup_owned_prompt || true
  record_run_finish "$final_rc"
  exit "$final_rc"
}
trap early_on_exit EXIT
trap 'exit 130' INT
trap 'exit 143' TERM

# Record before selection and preflight so early dispatch failures remain
# visible in the durable run ledger.
record_run_start
unset SUBFLEET_RUN_DECISION_JSON SUBFLEET_RUN_ORIGINAL_OUT SUBFLEET_RUN_LANE_LOG

pick_lane() {
  local args home
  if [ -n "${SUBFLEET_CODEX_PICK:-}" ]; then
    "$SUBFLEET_CODEX_PICK"
    return
  fi
  args=(pick codex)
  if [ "${#EXCLUDED_HOMES[@]}" -gt 0 ]; then
    for home in "${EXCLUDED_HOMES[@]}"; do
      args+=(--exclude "$home")
    done
  fi
  "$SUBFLEET" "${args[@]}"
}

if [ -z "$CODEX_HOME_ARG" ]; then
  CODEX_HOME_ARG=$(pick_lane) || {
    echo "subfleet codex: no dispatchable lane" >&2
    exit 1
  }
  [ -n "$CODEX_HOME_ARG" ] || {
    echo "subfleet codex: pick returned nothing" >&2
    exit 1
  }
  AUTO_LANE=1
  echo "subfleet codex: lane auto-assigned -> $CODEX_HOME_ARG" >&2
fi
record_run_lane

EXTRA_CFG=()
[ -n "$EFFORT" ] && EXTRA_CFG+=(-c "model_reasoning_effort=\"$EFFORT\"")
if [ "$SANDBOX" = "workspace-write" ]; then
  COMMON_DIR=$(git -C "$WORKDIR" rev-parse --path-format=absolute --git-common-dir 2>/dev/null || true)
  [ -n "$COMMON_DIR" ] && EXTRA_CFG+=(-c "sandbox_workspace_write.writable_roots=[\"$COMMON_DIR\"]")
fi
if [ "${SUBFLEET_CODEX_UNIFIED_EXEC:-on}" = "off" ]; then
  EXTRA_CFG+=(-c "features.unified_exec=false")
fi

CODEX_BIN=${SUBFLEET_CODEX_BIN:-}
if [ -z "$CODEX_BIN" ]; then
  CODEX_BIN=$("$SUBFLEET" _codex-binary 2>/dev/null) || CODEX_BIN=""
fi
[ -n "$CODEX_BIN" ] || {
  echo "subfleet codex: real codex binary not found" >&2
  exit 127
}

GUARD_ON=0
if [ "${SUBFLEET_CODEX_GUARD:-on}" = "off" ]; then
  echo "subfleet codex: command guard disabled" >&2
else
  GUARD_ON=1
  GUARD_OVERRIDE=$("$SCRIPT_DIR/subfleet-guard" override) || {
    echo "subfleet codex: guard override generation failed; refusing to launch (set SUBFLEET_CODEX_GUARD=off to bypass)" >&2
    exit 2
  }
  "$SCRIPT_DIR/subfleet-guard" preflight -H "$CODEX_HOME_ARG" -C "$WORKDIR" \
    --codex "$CODEX_BIN" >&2 || {
    echo "subfleet codex: guard preflight failed; refusing to launch (set SUBFLEET_CODEX_GUARD=off to bypass)" >&2
    exit 2
  }
  EXTRA_CFG+=(-c "$GUARD_OVERRIDE")
fi

salvage() {
  git -C "$WORKDIR" rev-parse --git-dir >/dev/null 2>&1 || return 0
  local salvage_sha=""
  if [ -n "$(git -C "$WORKDIR" status --porcelain 2>/dev/null)" ]; then
    local gitdir tmpidx tree parent branch ts ref
    gitdir=$(git -C "$WORKDIR" rev-parse --absolute-git-dir 2>/dev/null)
    parent=$(git -C "$WORKDIR" rev-parse --verify HEAD 2>/dev/null)
    if [ -n "$gitdir" ] && [ -n "$parent" ]; then
      tmpidx="$gitdir/codex-salvage-index.$$"
      if GIT_INDEX_FILE="$tmpidx" git -C "$WORKDIR" read-tree "$parent" 2>/dev/null \
         && GIT_INDEX_FILE="$tmpidx" git -C "$WORKDIR" add -A 2>/dev/null; then
        tree=$(GIT_INDEX_FILE="$tmpidx" git -C "$WORKDIR" write-tree 2>/dev/null)
        if [ -n "$tree" ]; then
          salvage_sha=$(git -C "$WORKDIR" commit-tree "$tree" -p "$parent" \
            -m "WIP salvage: subfleet codex auto-commit (model=$MODEL)" 2>/dev/null)
        fi
      fi
      rm -f "$tmpidx"
      if [ -n "$salvage_sha" ]; then
        branch=$(git -C "$WORKDIR" rev-parse --abbrev-ref HEAD 2>/dev/null || echo detached)
        ts=$(date +%Y%m%d-%H%M%S)
        ref="refs/codex-salvage/${branch//\//-}-$ts-$$"
        git -C "$WORKDIR" update-ref "$ref" "$salvage_sha" 2>/dev/null \
          && echo "subfleet codex: salvaged dirty state -> $ref ($(git -C "$WORKDIR" rev-parse --short "$salvage_sha")); branch/worktree untouched" >&2
      else
        echo "subfleet codex: WARNING dirty state present but salvage snapshot failed; worktree left as-is" >&2
      fi
    fi
  fi
  if [ -n "$SALVAGE_BRANCH" ]; then
    local push_src="${salvage_sha:-HEAD}"
    git -C "$WORKDIR" push -qf "$SALVAGE_REMOTE" "$push_src:refs/heads/$SALVAGE_BRANCH" 2>/dev/null \
      && echo "subfleet codex: pushed $push_src -> $SALVAGE_REMOTE/$SALVAGE_BRANCH" >&2
  fi
}

on_exit() {
  local final_rc=$?
  trap - EXIT
  if ! cleanup_owned_prompt; then
    echo "subfleet codex: WARNING skipping salvage and push to protect the prompt" >&2
  else
    salvage
  fi
  record_run_finish "$final_rc"
  exit "$final_rc"
}
trap on_exit EXIT

attempt=0
while :; do
  attempt=$((attempt + 1))
  : > "$OUT_FILE"
  CODEX_HOME="$CODEX_HOME_ARG" "$CODEX_BIN" exec --skip-git-repo-check --color never \
    --sandbox "$SANDBOX" ${EXTRA_CFG[@]+"${EXTRA_CFG[@]}"} -m "$MODEL" -C "$WORKDIR" -o "$OUT_FILE" \
    "$(cat "$PROMPT_FILE")" </dev/null >/dev/null 2>"$ERR_FILE"
  rc=$?
  if [ "$rc" -eq 0 ] && [ -s "$OUT_FILE" ]; then
    echo "subfleet codex: OK attempt=$attempt bytes=$(wc -c <"$OUT_FILE" | tr -d ' ')"
    exit 0
  fi

  if grep -qiE "content filter|trusted access|can('|’)t (help|assist) with" "$ERR_FILE"; then
    echo "subfleet codex: content filtered; rewrite the prompt before retrying" >&2
    exit 3
  fi
  if [ "$AUTO_LANE" = 1 ] && [ "$attempt" -le "$MAX_RETRIES" ] \
     && grep -qiE "hit your usage limit|usage limit reached" "$ERR_FILE"; then
    OLD_HOME=$CODEX_HOME_ARG
    EXCLUDED_HOMES+=("$OLD_HOME")
    NEW_HOME=$(pick_lane) || NEW_HOME=""
    if [ -n "$NEW_HOME" ] && [ "$NEW_HOME" != "$OLD_HOME" ]; then
      if [ "$GUARD_ON" = 1 ] \
         && ! "$SCRIPT_DIR/subfleet-guard" preflight -H "$NEW_HOME" -C "$WORKDIR" \
           --codex "$CODEX_BIN" >&2; then
        echo "subfleet codex: guard refused the re-picked lane $NEW_HOME" >&2
      else
        CODEX_HOME_ARG=$NEW_HOME
        record_run_lane
        echo "subfleet codex: $OLD_HOME hit its usage limit; re-picked -> $NEW_HOME" >&2
        continue
      fi
    else
      echo "subfleet codex: usage limit on $OLD_HOME and no alternate lane available" >&2
    fi
  fi
  if [ "$attempt" -le "$MAX_RETRIES" ] \
     && grep -qiE "capacity|401|unauthorized|rate limit|too many requests|disconnect|overloaded|timed? ?out|temporarily" "$ERR_FILE"; then
    wait_s=$((60 * attempt))
    echo "subfleet codex: transient failure (rc=$rc), retry $attempt/$MAX_RETRIES in ${wait_s}s" >&2
    sleep "$wait_s"
    continue
  fi
  echo "subfleet codex: FAILED rc=$rc attempt=$attempt; stderr tail:" >&2
  tail -3 "$ERR_FILE" >&2
  [ "$rc" -eq 0 ] && rc=1
  exit "$rc"
done
