#!/usr/bin/env bash
# ask-grok.sh — Claude -> Grok capability router.
#
# Routes a task to the Grok CLI (xAI Grok 4.3) for capabilities Claude lacks
# (image/video generation, live web + X/Twitter data) or for an independent
# second model (verification, drafts). Headless, safe-by-default, artifact-aware.
#
# Auth is Grok's own OIDC (~/.grok/auth.json, grok.com login) — NO API key here.
# Prompts are ALWAYS passed via --prompt-file (inline args break on shell quoting).
# Artifacts (image/video) are retrieved by globbing the session dir by sessionId,
# so they're recovered even when the headless run reports stopReason=Cancelled
# before echoing the path (a known finalization quirk).
#
# Usage:
#   ask-grok.sh text    "<prompt>"                         # text / reasoning / second opinion
#   ask-grok.sh verify  "<claim or finding to refute>"     # adversarial check (text + --check)
#   ask-grok.sh image   "<prompt>" [out.png]               # image_gen -> copy artifact to out
#   ask-grok.sh video   "<prompt>" <input-image> [out.mp4] # image_to_video on input image
#   ask-grok.sh xsearch "<query>"                          # live X/Twitter + web search
#   ask-grok.sh raw     <grok args...>                     # escape hatch (passes through)
#
# Exit: 0 ok, 1 no result/artifact, 2 usage, 127 grok not found.
set -uo pipefail
SELF="ask-grok"

PY="$(command -v py 2>/dev/null || command -v python 2>/dev/null || command -v python3 2>/dev/null || true)"
[ -z "$PY" ] && { echo "[$SELF] python (py/python/python3) required for JSON parsing" >&2; exit 127; }

# --- identity.json (per-machine, gitignored) declares whether grok is installed here ---
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" 2>/dev/null && pwd)"
IDFILE=""
[ -n "${CLAUDETOOLS_ROOT:-}" ] && [ -f "$CLAUDETOOLS_ROOT/.claude/identity.json" ] && IDFILE="$CLAUDETOOLS_ROOT/.claude/identity.json"
[ -z "$IDFILE" ] && IDFILE="$(cd "$SCRIPT_DIR/../../.." 2>/dev/null && pwd)/identity.json"
idgrok() {  # read field $1 from identity.json .grok (empty if absent)
  [ -f "$IDFILE" ] || { echo ""; return; }
  "$PY" -c "import json,sys
try:
    g=(json.load(sys.stdin).get('grok') or {}); v=g.get('$1','')
    print('' if v is None else (str(v).lower() if isinstance(v,bool) else v))
except Exception: print('')" < "$IDFILE"
}

# If identity explicitly says grok is NOT installed here, fail fast with routing guidance.
if [ "$(idgrok installed)" = "false" ]; then
  echo "[$SELF] grok is not installed on this machine (identity.json grok.installed=false)." >&2
  echo "[$SELF] Grok runs only on the fleet grok host. Route this request there (remote routing not yet wired) or install grok + set identity.json grok.installed=true." >&2
  exit 3
fi

# --- locate the grok binary: GROK env > identity.json grok.binary > auto-locate ---
GROK="${GROK:-}"
cand="$(idgrok binary)"
[ -z "$GROK" ] && [ -n "$cand" ] && [ -x "$cand" ] && GROK="$cand"
if [ -z "$GROK" ]; then
  if command -v grok >/dev/null 2>&1; then GROK="$(command -v grok)"; else
    for c in "$HOME/.grok/bin/grok.exe" "/c/Users/${USERNAME:-${USER:-x}}/.grok/bin/grok.exe" \
             "$HOME/.grok/bin/grok" "${LOCALAPPDATA:-}/Programs/grok/grok.exe"; do
      [ -x "$c" ] && { GROK="$c"; break; }
    done
  fi
fi
[ -z "$GROK" ] && { echo "[$SELF] grok CLI not found (set identity.json grok.binary, GROK=, or install grok)" >&2; exit 127; }

MODE="${1:-}"; shift 2>/dev/null || true
[ -z "$MODE" ] && { echo "usage: $SELF {text|verify|image|video|xsearch|raw} ..." >&2; exit 2; }

TMP="$(mktemp -d)"; trap 'rm -rf "$TMP"' EXIT
WORK="$TMP/work"; mkdir -p "$WORK"
PF="$TMP/prompt.txt"; OUT="$TMP/out.json"
RUN_CWD="$WORK"   # grok's working dir; the 'review' mode overrides to the repo so read_file can reach repo files
REPO_ROOT="${CLAUDETOOLS_ROOT:-$(cd "$SCRIPT_DIR/../../../.." 2>/dev/null && pwd)}"

# run grok headless. $1=timeout secs; rest=extra flags. Reads $PF -> $OUT.
# Never fails the script on grok's exit code (Cancelled is expected; we read artifacts).
# Use gtimeout on macOS (from brew coreutils), timeout on Linux/Windows.
TIMEOUT_CMD="timeout"
if [[ "$OSTYPE" == "darwin"* ]]; then
  TIMEOUT_CMD="$(command -v gtimeout 2>/dev/null || echo timeout)"
fi

run_grok() {
  local to="$1"; shift
  "$TIMEOUT_CMD" "$to" "$GROK" --prompt-file "$PF" --output-format json \
    --permission-mode dontAsk --no-subagents --no-plan --cwd "$RUN_CWD" "$@" \
    >"$OUT" 2>"$TMP/err.txt" || true
}

# parse a top-level field from $OUT. Read via stdin so Windows python never has
# to resolve the git-bash (/c/...) path itself.
jfield() { "$PY" -c "import json,sys
try:
    d=json.load(sys.stdin); print(d.get('$1','') or '')
except Exception:
    print('')" < "$OUT"; }

# newest artifact under any session dir for this sessionId: $1=sid $2=images|videos
find_artifact() {
  ls -t "$HOME/.grok/sessions/"*"/$1/$2/"* 2>/dev/null | head -1
}

case "$MODE" in
  text|verify)
    # content from --prompt-file <path> (good for long docs) or the positional arg
    SRC=""
    if [ "${1:-}" = "--prompt-file" ]; then
      [ -f "${2:-}" ] || { echo "[$SELF] prompt file not found: ${2:-}" >&2; exit 2; }
      SRC="$(cat "$2")"
    else
      SRC="${1:-}"
    fi
    [ -z "$SRC" ] && { echo "usage: $SELF $MODE \"<prompt>\" | $SELF $MODE --prompt-file <path>" >&2; exit 2; }
    # Prompt-level steering keeps it text-only and (for verify) adversarial.
    # (The --disallowed-tools/--rules flags tripped the CLI; --check adds a slow
    # multi-turn self-check loop — both avoided in favor of prompt steering.)
    if [ "$MODE" = "verify" ]; then
      printf 'Adversarially evaluate the following claim/finding/document: try hard to find any reason it is WRONG, incomplete, or overstated. Then give a verdict plus specific justification. Answer in text only; do not use tools. Content:\n%s' "$SRC" > "$PF"
    else
      printf 'Answer directly in text; do not use tools.\n%s' "$SRC" > "$PF"
    fi
    run_grok 180 --disable-web-search --max-turns 3
    txt="$(jfield text)"
    if [ -n "$txt" ]; then printf '%s\n' "$txt"; else
      echo "[$SELF] no text (stopReason=$(jfield stopReason)); raw: $OUT" >&2; exit 1; fi
    ;;
  image)
    [ -z "${1:-}" ] && { echo "usage: $SELF image \"<prompt>\" [out.png]" >&2; exit 2; }
    out="${2:-grok-image.png}"
    printf 'Use your image_gen tool to generate exactly this image, save it, then stop. Image: %s' "$1" > "$PF"
    run_grok 240 --disable-web-search --max-turns 12
    sid="$(jfield sessionId)"; art="$(find_artifact "$sid" images)"
    if [ -n "$art" ] && [ -f "$art" ]; then cp -f "$art" "$out"
      echo "[$SELF] image OK -> $out  (session $sid)"
    else echo "[$SELF] no image artifact (session=$sid, stopReason=$(jfield stopReason))" >&2; exit 1; fi
    ;;
  video)
    [ -z "${1:-}" ] || [ -z "${2:-}" ] && { echo "usage: $SELF video \"<prompt>\" <input-image> [out.mp4]" >&2; exit 2; }
    input="$2"; out="${3:-grok-video.mp4}"
    [ -f "$input" ] || { echo "[$SELF] input image not found: $input" >&2; exit 2; }
    cp -f "$input" "$WORK/input.jpg"
    printf 'Use your image_to_video tool to animate input.jpg (in the current directory) into a short clip, save it, then stop. Motion: %s' "$1" > "$PF"
    run_grok 360 --disable-web-search --max-turns 20
    sid="$(jfield sessionId)"; art="$(find_artifact "$sid" videos)"
    if [ -n "$art" ] && [ -f "$art" ]; then cp -f "$art" "$out"
      echo "[$SELF] video OK -> $out  (session $sid)"
    else echo "[$SELF] no video artifact (session=$sid, stopReason=$(jfield stopReason))" >&2; exit 1; fi
    ;;
  xsearch)
    [ -z "${1:-}" ] && { echo "usage: $SELF xsearch \"<query>\"" >&2; exit 2; }
    printf 'Use your web_search and X/Twitter search tools to answer this, cite sources, then stop: %s' "$1" > "$PF"
    run_grok 150 --max-turns 6
    txt="$(jfield text)"
    if [ -n "$txt" ]; then printf '%s\n' "$txt"; else
      echo "[$SELF] no result (stopReason=$(jfield stopReason))" >&2; exit 1; fi
    ;;
  review|file)
    [ -z "${1:-}" ] && { echo "usage: $SELF review <file-path> [instructions]" >&2; exit 2; }
    target="$1"
    instr="${2:-Give an independent, critical review of this file: accuracy, gaps/omissions, and concrete improvements. Be specific.}"
    # Grok reads the file itself (no embedding) -- run it in the repo so read_file resolves repo-relative paths.
    [ -f "$target" ] || [ -f "$REPO_ROOT/$target" ] || { echo "[$SELF] file not found: $target" >&2; exit 2; }
    RUN_CWD="$REPO_ROOT"
    printf 'Use your read_file tool to read the file at this path (relative to your current directory), then do the task and stop. You may also read closely-related files it references if that helps. Do not modify anything.\nPath: %s\n\nTask: %s' "$target" "$instr" > "$PF"
    run_grok 240 --max-turns 12
    txt="$(jfield text)"
    if [ -n "$txt" ]; then printf '%s\n' "$txt"; else
      echo "[$SELF] no result (session=$(jfield sessionId), stopReason=$(jfield stopReason))" >&2; exit 1; fi
    ;;
  raw)
    "$GROK" "$@"
    ;;
  *)
    echo "[$SELF] unknown mode '$MODE' (use text|verify|image|video|xsearch|raw)" >&2; exit 2 ;;
esac