#!/usr/bin/env bash
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
PASS=0
FAIL=0

ok() { PASS=$((PASS+1)); printf '  ok   %s\n' "$1"; }
bad() { FAIL=$((FAIL+1)); printf '  FAIL %s\n     %s\n' "$1" "$2"; }

mkfixture() {
  local dir
  dir="$(mktemp -d "/tmp/grok-session-XXXX")"
  mkdir -p "$dir/wrappers" "$dir/bin" "$dir/home"
  cp "$ROOT/wrappers/grok.sh" "$dir/wrappers/grok.sh"
  mkdir -p "$dir/wrappers/lib"
  cp "$ROOT/wrappers/lib/verdict-status.py" "$dir/wrappers/lib/verdict-status.py"
  chmod +x "$dir/wrappers/grok.sh"
  printf '%s\n' "$dir"
}

write_mock_grok() {
  local dir="$1"
  cat > "$dir/bin/grok" <<'EOF'
#!/usr/bin/env bash
set -euo pipefail

log_file="${MOCK_LOG_FILE:?}"
printf '%s\n' "$*" >> "$log_file"

if [[ "${1:-}" == "version" || "${1:-}" == "--version" ]]; then
  printf '%s\n' "${MOCK_VERSION_OUTPUT:-grok 0.0.0-mock}"
  exit "${MOCK_VERSION_EXIT_CODE:-0}"
fi

if [[ "${1:-}" == "models" ]]; then
  printf '%s\n' "${MOCK_MODELS_OUTPUT:-grok-4.5}"
  exit "${MOCK_MODELS_EXIT_CODE:-0}"
fi

# Headless dispatch: echo run output then exit
if [[ -n "${MOCK_RUN_OUTPUT+x}" ]]; then
  printf '%s\n' "$MOCK_RUN_OUTPUT"
else
  printf '%s\n' '{"type":"text","data":"ok"}'
  printf '%s\n' '{"type":"end","sessionId":"sess-from-stream","stopReason":"EndTurn"}'
fi
exit "${MOCK_RUN_EXIT_CODE:-0}"
EOF
  chmod +x "$dir/bin/grok"
}

run_missing_model_case() {
  local dir out rc
  dir="$(mkfixture)"
  write_mock_grok "$dir"
  if out=$(
    PATH="$dir/bin:$PATH" \
    _GROK_ENGINE_BIN="$dir/bin/grok" \
    MOCK_LOG_FILE="$dir/invocations.log" \
    HOME="$dir/home" \
    "$dir/wrappers/grok.sh" \
      --workspace "$dir" \
      --trust "hi" \
      --task-slug "t1" \
      --timeout 10 2>&1
  ); then
    rc=0
  else
    rc=$?
  fi
  if [[ "$rc" -eq 2 ]]; then
    ok "missing --model exits 2"
  else
    bad "missing --model exits 2" "rc=$rc out=${out:-<empty>}"
  fi
  rm -rf "$dir"
}

run_dispatch_case() {
  local dir out rc log_file status_line
  dir="$(mkfixture)"
  write_mock_grok "$dir"
  log_file="$dir/invocations.log"
  : > "$log_file"

  if out=$(
    PATH="$dir/bin:$PATH" \
    _GROK_ENGINE_BIN="$dir/bin/grok" \
    MOCK_LOG_FILE="$log_file" \
    HOME="$dir/home" \
    "$dir/wrappers/grok.sh" \
      --workspace "$dir" \
      --trust "ship it" \
      --task-slug "task-1" \
      --model "grok-composer-2.5-fast" \
      --timeout 30 2>&1
  ); then
    rc=0
  else
    rc=$?
  fi

  if [[ "$rc" -eq 0 ]]; then
    ok "dispatch exits zero"
  else
    bad "dispatch exits zero" "rc=$rc out=${out:-<empty>}"
  fi

  # Last non-version invocation is the run line
  local run_line
  run_line="$(grep -E '^-p |^-p$' "$log_file" | tail -n 1 || true)"
  # mock logs full argv as one line; match key flags
  if grep -Fq -- '-p ship it' "$log_file" || grep -Fq -- $'-p\nship it' "$log_file"; then
    :
  fi
  if grep -E -- '-p ship it|-p "ship it"' "$log_file" >/dev/null || grep -Fq -- '-p ship it' "$log_file"; then
    ok "dispatch passes -p prompt"
  elif grep -Fq 'ship it' "$log_file" && grep -Fq -- '-p' "$log_file"; then
    ok "dispatch passes -p prompt"
  else
    bad "dispatch passes -p prompt" "$(cat "$log_file")"
  fi

  if grep -Fq -- '-m grok-composer-2.5-fast' "$log_file"; then
    ok "dispatch pins model"
  else
    bad "dispatch pins model" "$(cat "$log_file")"
  fi

  if grep -Fq -- "--cwd $dir" "$log_file"; then
    ok "dispatch pins --cwd workspace"
  else
    bad "dispatch pins --cwd workspace" "$(cat "$log_file")"
  fi

  if grep -Fq -- '--output-format streaming-json' "$log_file"; then
    ok "dispatch uses streaming-json"
  else
    bad "dispatch uses streaming-json" "$(cat "$log_file")"
  fi

  if grep -Fq -- '--always-approve' "$log_file"; then
    ok "dispatch uses --always-approve"
  else
    bad "dispatch uses --always-approve" "$(cat "$log_file")"
  fi

  status_line="$(printf '%s\n' "$out" | grep -E '^\{\"ok\":' | tail -n 1 || true)"
  if printf '%s\n' "$status_line" | grep -Fq '"ok":true' && printf '%s\n' "$status_line" | grep -Fq 'sess-from-stream'; then
    ok "dispatch prints ok status with session_id from stream"
  else
    bad "dispatch prints ok status with session_id from stream" "${status_line:-<empty>} full=${out:-}"
  fi

  if printf '%s\n' "$out" | grep -Fq '"kind":"usage"'; then
    ok "dispatch emits usage event"
  else
    bad "dispatch emits usage event" "${out:-<empty>}"
  fi

  rm -rf "$dir"
}

run_resume_case() {
  local dir out rc log_file
  dir="$(mkfixture)"
  write_mock_grok "$dir"
  log_file="$dir/invocations.log"
  : > "$log_file"

  if out=$(
    PATH="$dir/bin:$PATH" \
    _GROK_ENGINE_BIN="$dir/bin/grok" \
    MOCK_LOG_FILE="$log_file" \
    HOME="$dir/home" \
    "$dir/wrappers/grok.sh" \
      --workspace "$dir" \
      --trust "resume it" \
      --task-slug "task-2" \
      --model "grok-4.5" \
      --resume "sess-existing-9" \
      --timeout 30 2>&1
  ); then
    rc=0
  else
    rc=$?
  fi

  if [[ "$rc" -eq 0 ]]; then
    ok "resume dispatch exits zero"
  else
    bad "resume dispatch exits zero" "rc=$rc out=${out:-<empty>}"
  fi

  if grep -Fq -- '--resume sess-existing-9' "$log_file"; then
    ok "resume dispatch forwards session id"
  else
    bad "resume dispatch forwards session id" "$(cat "$log_file")"
  fi

  rm -rf "$dir"
}

run_health_case() {
  local dir out rc log_file
  dir="$(mkfixture)"
  write_mock_grok "$dir"
  log_file="$dir/invocations.log"
  : > "$log_file"

  if out=$(
    PATH="$dir/bin:$PATH" \
    _GROK_ENGINE_BIN="$dir/bin/grok" \
    MOCK_LOG_FILE="$log_file" \
    MOCK_VERSION_OUTPUT='grok 9.9.9-test' \
    HOME="$dir/home" \
    "$dir/wrappers/grok.sh" --health 2>&1
  ); then
    rc=0
  else
    rc=$?
  fi

  if [[ "$rc" -eq 0 ]] && printf '%s\n' "$out" | grep -Fq '"ok":true' && printf '%s\n' "$out" | grep -Fq '9.9.9-test'; then
    ok "health mode returns version json"
  else
    bad "health mode returns version json" "rc=$rc out=${out:-<empty>}"
  fi

  if grep -E '^(version|--version)$' "$log_file" >/dev/null || grep -Fq 'version' "$log_file"; then
    ok "health mode probes version"
  else
    bad "health mode probes version" "$(cat "$log_file")"
  fi

  rm -rf "$dir"
}

run_list_models_case() {
  local dir out rc log_file
  dir="$(mkfixture)"
  write_mock_grok "$dir"
  log_file="$dir/invocations.log"
  : > "$log_file"

  if out=$(
    PATH="$dir/bin:$PATH" \
    _GROK_ENGINE_BIN="$dir/bin/grok" \
    MOCK_LOG_FILE="$log_file" \
    MOCK_MODELS_OUTPUT=$'grok-4.5\ngrok-composer-2.5-fast' \
    HOME="$dir/home" \
    "$dir/wrappers/grok.sh" --list-models 2>&1
  ); then
    rc=0
  else
    rc=$?
  fi

  if [[ "$rc" -eq 0 ]] && [[ "$out" == $'grok-4.5\ngrok-composer-2.5-fast' ]]; then
    ok "list-models proxies model listing"
  else
    bad "list-models proxies model listing" "rc=$rc out=${out:-<empty>}"
  fi

  if grep -Fxq 'models' "$log_file"; then
    ok "list-models calls grok models"
  else
    bad "list-models calls grok models" "$(cat "$log_file")"
  fi

  rm -rf "$dir"
}

run_timeout_case() {
  local dir out rc
  dir="$(mkfixture)"
  write_mock_grok "$dir"

  if out=$(
    PATH="$dir/bin:$PATH" \
    _GROK_ENGINE_BIN="$dir/bin/grok" \
    MOCK_LOG_FILE="$dir/invocations.log" \
    MOCK_RUN_EXIT_CODE=124 \
    HOME="$dir/home" \
    "$dir/wrappers/grok.sh" \
      --workspace "$dir" \
      --trust "slow" \
      --task-slug "timeout-1" \
      --model "grok-4.5" \
      --timeout 30 2>&1
  ); then
    rc=0
  else
    rc=$?
  fi

  if [[ "$rc" -eq 124 ]]; then
    ok "timeout maps to contract 124"
  else
    bad "timeout maps to contract 124" "rc=$rc out=${out:-<empty>}"
  fi
  rm -rf "$dir"
}

run_rate_limit_case() {
  local dir out rc
  dir="$(mkfixture)"
  write_mock_grok "$dir"

  if out=$(
    PATH="$dir/bin:$PATH" \
    _GROK_ENGINE_BIN="$dir/bin/grok" \
    MOCK_LOG_FILE="$dir/invocations.log" \
    MOCK_RUN_EXIT_CODE=1 \
    MOCK_RUN_OUTPUT='rate limit exceeded try again later' \
    HOME="$dir/home" \
    "$dir/wrappers/grok.sh" \
      --workspace "$dir" \
      --trust "quota" \
      --task-slug "rl-1" \
      --model "grok-4.5" \
      --timeout 30 2>&1
  ); then
    rc=0
  else
    rc=$?
  fi

  if [[ "$rc" -eq 75 ]]; then
    ok "rate-limit maps to contract 75"
  else
    bad "rate-limit maps to contract 75" "rc=$rc out=${out:-<empty>}"
  fi

  if printf '%s\n' "$out" | grep -Fq 'errorKind":"quota"'; then
    ok "rate-limit emits quota error event"
  else
    bad "rate-limit emits quota error event" "${out:-<empty>}"
  fi
  rm -rf "$dir"
}

run_transcript_case() {
  local dir out rc transcript
  dir="$(mkfixture)"
  write_mock_grok "$dir"
  transcript="$dir/transcript.jsonl"

  if out=$(
    PATH="$dir/bin:$PATH" \
    _GROK_ENGINE_BIN="$dir/bin/grok" \
    MOCK_LOG_FILE="$dir/invocations.log" \
    HARNESS_TRANSCRIPT_PATH="$transcript" \
    HOME="$dir/home" \
    "$dir/wrappers/grok.sh" \
      --workspace "$dir" \
      --trust "tee me" \
      --task-slug "tx-1" \
      --model "grok-4.5" \
      --timeout 30 2>&1
  ); then
    rc=0
  else
    rc=$?
  fi

  if [[ "$rc" -eq 0 ]] && [[ -s "$transcript" ]]; then
    ok "transcript path receives stream"
  else
    bad "transcript path receives stream" "rc=$rc size=$(wc -c <"$transcript" 2>/dev/null || echo 0)"
  fi
  rm -rf "$dir"
}

run_probe_fail_case() {
  local dir out rc
  dir="$(mkfixture)"
  write_mock_grok "$dir"

  if out=$(
    PATH="$dir/bin:$PATH" \
    _GROK_ENGINE_BIN="$dir/bin/grok" \
    MOCK_LOG_FILE="$dir/invocations.log" \
    MOCK_VERSION_EXIT_CODE=1 \
    HOME="$dir/home" \
    "$dir/wrappers/grok.sh" \
      --workspace "$dir" \
      --trust "nope" \
      --task-slug "down-1" \
      --model "grok-4.5" \
      --timeout 30 2>&1
  ); then
    rc=0
  else
    rc=$?
  fi

  if [[ "$rc" -eq 3 ]]; then
    ok "version probe failure exits 3 without dispatch"
  else
    bad "version probe failure exits 3 without dispatch" "rc=$rc out=${out:-<empty>}"
  fi
  rm -rf "$dir"
}

echo "grok session wrapper:"
run_missing_model_case
run_dispatch_case
run_resume_case
run_health_case
run_list_models_case
run_timeout_case
run_rate_limit_case
run_transcript_case
run_probe_fail_case
echo
echo "PASS=$PASS FAIL=$FAIL"
[[ "$FAIL" -eq 0 ]]