#!/usr/bin/env bash
# grok.sh — dispatch Grok Build headless (grok -p) under WRAPPER-CONTRACT.
# audience: AI coding agents first. Invoke by path; do NOT re-derive the run line.
# Exit: 0 = completion, 124 = timeout/killed, 2 = usage error, 75 = rate-limited, 3 = engine down OR agent
#   returned 0 with empty captured output (silent write failure / disk full — success requires evidence).
set -uo pipefail

WORKSPACE="" PROMPT="" TASK_SLUG="" MODEL="" RESUME_ID="" PROFILE="" TIMEOUT=360
MODE="dispatch"
while [[ $# -gt 0 ]]; do
  case "$1" in
    --workspace) [[ $# -lt 2 ]] && { echo '{"ok":false,"detail":"--workspace requires a value"}' >&2; exit 2; }; WORKSPACE="${2:-}"; shift 2;;
    --trust) [[ $# -lt 2 ]] && { echo '{"ok":false,"detail":"--trust requires a value"}' >&2; exit 2; }; PROMPT="${2:-}"; shift 2;;
    --task-slug) [[ $# -lt 2 ]] && { echo '{"ok":false,"detail":"--task-slug requires a value"}' >&2; exit 2; }; TASK_SLUG="${2:-}"; shift 2;;
    --model) [[ $# -lt 2 ]] && { echo '{"ok":false,"detail":"--model requires a value"}' >&2; exit 2; }; MODEL="${2:-}"; shift 2;;
    --resume) [[ $# -lt 2 ]] && { echo '{"ok":false,"detail":"--resume requires a session id"}' >&2; exit 2; }; RESUME_ID="${2:-}"; shift 2;;
    --timeout) [[ $# -lt 2 ]] && { echo '{"ok":false,"detail":"--timeout requires a value"}' >&2; exit 2; }; TIMEOUT="${2:-}"; shift 2;;
    --profile) [[ $# -lt 2 ]] && { echo '{"ok":false,"detail":"--profile requires a value"}' >&2; exit 2; }; PROFILE="${2:-}"; shift 2;;
    --health) MODE="health"; shift;;
    --list-models) MODE="list-models"; shift;;
    *) shift;;
  esac
done
[[ -n "$PROFILE" ]] && echo "notice: --profile=$PROFILE ignored — grok.sh has no account routing" >&2

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VERDICT_STATUS_PY="$SCRIPT_DIR/lib/verdict-status.py"
[[ -f "$VERDICT_STATUS_PY" ]] || { echo '{"ok":false,"detail":"verdict-status.py missing — engine precondition unmet"}' >&2; exit 3; }
source "$SCRIPT_DIR/lib/finalize.sh"
source "$SCRIPT_DIR/lib/spend-cap.sh"
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
_GROK_BIN="${_GROK_ENGINE_BIN:-grok}"

load_env() {
  local env_candidates=(
    "${_GROK_ENV:-}"
    "$REPO_ROOT/GROK.env"
    "$HOME/.claude/workflows/lib/GROK.env"
  )
  local env_file
  for env_file in "${env_candidates[@]}"; do
    [[ -n "$env_file" && -f "$env_file" ]] || continue
    source "$env_file" || return 1
    return 0
  done
  return 0
}

make_log_dir() {
  local dir="$HOME/Projects/mega-plan-harness/tmp/logs"
  if ! mkdir -p "$dir" 2>/dev/null || [[ ! -d "$dir" || ! -w "$dir" ]]; then
    dir="${TMPDIR:-/tmp}/mega-plan-harness-logs"
    mkdir -p "$dir" || return 1
  fi
  printf '%s\n' "$dir"
}

probe_version() {
  local out
  if out="$("$_GROK_BIN" version 2>/dev/null)"; then
    printf '%s\n' "$out" | tail -n 1 | tr -d '\r'
    return 0
  fi
  if out="$("$_GROK_BIN" --version 2>/dev/null)"; then
    printf '%s\n' "$out" | tail -n 1 | tr -d '\r'
    return 0
  fi
  return 1
}

prepare_grok_engine() {
  if ! load_env; then
    echo '{"ok":false,"detail":"GROK.env source failed — engine precondition unmet"}' >&2
    return 3
  fi
  if ! command -v "$_GROK_BIN" >/dev/null 2>&1 && [[ ! -x "$_GROK_BIN" ]]; then
    echo '{"ok":false,"detail":"grok binary not found"}' >&2
    return 3
  fi
}

case "$MODE" in
  health)
    prepare_grok_engine || exit $?
    version_out="$(probe_version)" || { echo '{"ok":false,"detail":"grok version probe failed"}' >&2; exit 3; }
    VERSION_OUT="$version_out" python3 - <<'PY'
import json
import os
print(json.dumps({"ok": True, "version": os.environ.get("VERSION_OUT", "")}, separators=(",", ":")))
PY
    exit 0
    ;;
  list-models)
    prepare_grok_engine || exit $?
    "$_GROK_BIN" models
    exit $?
    ;;
esac

TASK_SLUG="${TASK_SLUG//[^a-zA-Z0-9_-]/}"
RUN_ID="$(date +%Y-%m-%d-%H-%M-%S)-$$"

[[ -n "$WORKSPACE" && -d "$WORKSPACE" ]] || { echo '{"ok":false,"detail":"--workspace missing or not a dir"}' >&2; exit 2; }
[[ -n "$PROMPT" ]] || { echo '{"ok":false,"detail":"--trust (prompt) required"}' >&2; exit 2; }
[[ -n "$TASK_SLUG" ]] || { echo '{"ok":false,"detail":"--task-slug required (log + retry hygiene)"}' >&2; exit 2; }
[[ -n "$MODEL" ]] || { echo '{"ok":false,"detail":"--model required (no silent default)"}' >&2; exit 2; }

spend_cap_refuse_if_capped roy-grok 'grok' "$RESUME_ID" '' "$MODEL"
prepare_grok_engine || exit $?

if ! probe_version >/dev/null; then
  echo '{"ok":false,"detail":"grok version probe failed — engine precondition unmet"}' >&2
  exit 3
fi

LOG_DIR="$(make_log_dir)" || { echo '{"ok":false,"detail":"log dir create failed"}' >&2; exit 3; }
LOG="$LOG_DIR/$RUN_ID-$TASK_SLUG.log"
RAW_OUTPUT_FILE="$LOG_DIR/$RUN_ID-$TASK_SLUG.raw"
TRANSCRIPT_PATH="${HARNESS_TRANSCRIPT_PATH:-}"
if [[ -n "$TRANSCRIPT_PATH" ]]; then
  mkdir -p "$(dirname "$TRANSCRIPT_PATH")" || { echo '{"ok":false,"detail":"transcript dir create failed"}' >&2; exit 3; }
fi

stream_output() {
  if [[ -n "$TRANSCRIPT_PATH" ]]; then
    tee -a "$LOG" "$RAW_OUTPUT_FILE" "$TRANSCRIPT_PATH"
  else
    tee -a "$LOG" "$RAW_OUTPUT_FILE"
  fi
}

emit_event_line() {
  local line="$1"
  if [[ -n "$TRANSCRIPT_PATH" ]]; then
    printf '%s\n' "$line" | tee -a "$LOG" "$TRANSCRIPT_PATH"
  else
    printf '%s\n' "$line" | tee -a "$LOG"
  fi
}

usage_event_json() {
  python3 - "$MODEL" <<'PY'
import json
import sys
print(json.dumps({"kind": "usage", "usage": {"unknown": True, "model": sys.argv[1]}}, separators=(",", ":")))
PY
}

# Build run args once — same array for log header and dispatch (no divergence).
GROK_RUN=("-p" "$PROMPT" "-m" "$MODEL" "--cwd" "$WORKSPACE" "--output-format" "streaming-json" "--always-approve")
[[ -n "$RESUME_ID" ]] && GROK_RUN+=(--resume "$RESUME_ID")
GROK_ARGS=(timeout -k 5 "$TIMEOUT" "$_GROK_BIN" "${GROK_RUN[@]}")
GROK_CMD=$(printf '%q ' "${GROK_ARGS[@]}")
{
  echo "=== TOOL: grok (${MODEL}) ==="
  echo "=== CMD: $GROK_CMD ==="
  echo "=== workspace: $WORKSPACE  timeout: ${TIMEOUT}s  log: $LOG ==="
  [[ -n "$RESUME_ID" ]] && echo "=== resume: $RESUME_ID ==="
} | tee "$LOG"

cd "$WORKSPACE" || { echo '{"ok":false,"detail":"cd workspace failed"}' >&2; exit 2; }
"${GROK_ARGS[@]}" < /dev/null 2>&1 | stream_output
RC="${PIPESTATUS[0]}"
raw_output="$(cat "$RAW_OUTPUT_FILE" 2>/dev/null || true)"

if [[ $RC -eq 0 ]]; then
  final_rc=0
elif [[ $RC -eq 124 || $RC -eq 137 ]]; then
  final_rc=124
elif [[ $RC -eq 2 ]]; then
  final_rc=2
elif grep -qiE 'usage limit|rate.?limit|too many requests|(^|[^0-9])429([^0-9]|$)' <<<"$raw_output"; then
  final_rc=75
else
  final_rc=3
fi
if [[ $final_rc -eq 0 ]] && ! [[ "$raw_output" =~ [^[:space:]] ]]; then final_rc=3; fi

echo "=== grok.sh exit=$RC (mapped to contract exit $final_rc) ===" | tee -a "$LOG"

SESSION_ID="$(
  RAW_OUTPUT_FILE="$RAW_OUTPUT_FILE" EXPLICIT_RESUME="$RESUME_ID" python3 - <<'PY'
import json
import os

path = os.environ.get("RAW_OUTPUT_FILE", "")
session_id = os.environ.get("EXPLICIT_RESUME", "") or ""
try:
    with open(path, "r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line.startswith("{"):
                continue
            try:
                parsed = json.loads(line)
            except Exception:
                continue
            if not isinstance(parsed, dict):
                continue
            value = parsed.get("sessionId") or parsed.get("session_id")
            if isinstance(value, str) and value:
                session_id = value
except Exception:
    pass
print(session_id)
PY
)"

status_json() {
  local line
  line="$(RAW_OUTPUT_FILE="$RAW_OUTPUT_FILE" OK="$1" DETAIL="$2" THREAD_ID="$3" python3 "$VERDICT_STATUS_PY")" || return 1
  [[ -n "${line//[[:space:]]/}" ]] || return 1
  printf '%s\n' "$line"
}

if [[ $final_rc -eq 75 ]]; then
  emit_event_line '{"kind":"error","subtype":"rate_limited","error":{"code":"rate_limit","message":"grok quota exhausted","errorKind":"quota"}}'
  STATUS_JSON="$(printf '{"ok":false,"detail":"rate-limited","thread_id":"%s"}' "$SESSION_ID")"
elif [[ $final_rc -eq 0 ]]; then
  STATUS_JSON="$(status_json 1 'grok completed' "$SESSION_ID")" || {
    echo '{"ok":false,"detail":"verdict-status.py failed — cannot honor status-line contract"}' >&2
    exit 3
  }
else
  STATUS_JSON="$(status_json 0 'grok failed' "$SESSION_ID")" || STATUS_JSON="$(printf '{"ok":false,"detail":"grok failed","thread_id":"%s"}' "$SESSION_ID")"
fi
emit_event_line "$(usage_event_json)"
printf '%s\n' "$STATUS_JSON" | tee -a "$LOG"
exit "$final_rc"