#!/usr/bin/env bash
# tests/lib.sh -- tiny assertion helpers for the ai-sandbox test suite.
#
# Every test file gets: AISB (path to the script), WORK (scratch dir), and uses
# sb() to run the sandbox with fully redirected XDG dirs so nothing touches the
# real machine.

PASS=0
FAIL=0
FAILED_NAMES=()
_current=""

t_name() { _current="$1"; printf '\n\033[1m# %s\033[0m\n' "$1"; }

_ok()   { PASS=$((PASS + 1)); printf '  \033[32mok\033[0m   %s\n' "$1"; }
_bad()  { FAIL=$((FAIL + 1)); FAILED_NAMES+=("$_current :: $1"); printf '  \033[31mFAIL\033[0m %s\n' "$1"; shift || true;
          [[ -n "${1:-}" ]] && printf '       %s\n' "$1"; return 0; }

# sb [args...] -> sets RC, OUT (stdout+stderr)
sb() {
  OUT="$("$AISB" "$@" 2>&1)" && RC=0 || RC=$?
  return 0
}

# sb_timeout <seconds> [args...] -- like sb, but bounded. Use it when a test expects
# the script to refuse: the refusal happens before podman, and the timeout is there in
# case a regression turns "refused instantly" into "hung forever".
sb_timeout() {
  local t="$1"; shift
  OUT="$(timeout "$t" "$AISB" "$@" 2>&1 </dev/null)" && RC=$? || RC=$?
  return 0
}

# Extract one function from the script itself, the way 06_entrypoint.sh extracts the
# entrypoint: the test then asks the real code a question instead of restating its answer.
aisb_fn() { # aisb_fn <function-name> -> that function's source, on stdout
  awk -v fn="$1" '$0 ~ "^" fn "\\(\\) \\{" {f=1} f {print} f && /^\}$/ {exit}' "$AISB"
}

# Where the sandbox would mount a host path. Empty means it mirrors the path; anything
# else names the container path the bind would bury. Extracted rather than reimplemented:
# the answer depends on the container's own filesystem (/usr, /tmp, /home/sandbox), and
# most of that hierarchy cannot be created by a test.
mirror_block_of() { # mirror_block_of <absolute host path> -> blocking path or empty
  local src; src="$(aisb_fn mirror_block_of)"
  [[ -n "$src" ]] || { echo "mirror_block_of not found in $AISB" >&2; return 2; }
  ( set -uo pipefail; eval "$src"; mirror_block_of "$1" )
}

# A directory whose host path the sandbox is allowed to mirror, in USER_MIRROR_DIR (empty
# when there is none). The scratch dir ($WORK) is under $TMPDIR, and /tmp is one of the
# paths the sandbox owns, so a run from there always takes the fallback branch -- the
# mirroring branch needs a per-user tree, and $HOME is the one place an unprivileged test
# can create it. Nothing is written into the real ~/.config, ~/.local or ~/.cache.
#
# Assigns a global instead of printing: the cleanup trap below belongs to the test file,
# so this must not be called as $(mirrorable_dir) -- a command substitution would set the
# trap in a subshell that exits, and takes the directory with it.
USER_MIRROR_DIR=""
mirrorable_dir() {
  local base="${HOME%/}" qbase
  USER_MIRROR_DIR=""
  [[ -n "$base" && -d "$base" && -w "$base" ]] || return 1
  [[ -z "$(mirror_block_of "$base")" ]] || return 1
  USER_MIRROR_DIR="$(mktemp -d "$base/.ai-sandbox-test-mirror.XXXXXX" 2>/dev/null)" || return 1
  [[ -n "$USER_MIRROR_DIR" ]] || return 1
  # Removed when the test file's subshell exits, including on a failed assert. The base is
  # quoted in now rather than read from $HOME later, so the trap cannot widen.
  printf -v qbase '%q' "$base"
  # shellcheck disable=SC2064   # immediate expansion of the captured base is the point
  trap "rm -rf -- $qbase/.ai-sandbox-test-mirror.* >/dev/null 2>&1 || true" EXIT
  return 0
}

# Run a command with stdin, stdout and stderr on a pty: -t and -i are decided from the
# real terminal, so a test about them cannot use the runner's pipes. util-linux's 'script'
# is the usual tool; python3's pty module does the same job where script is not installed
# (measured: all three descriptors are ttys and the exit code propagates). Sets OUT and RC
# like sb(). Args are joined for 'script', so pass simple words, not quoted sentences.
pty_available() { command -v script > /dev/null 2>&1 || command -v python3 > /dev/null 2>&1; }
pty_run() {
  OUT="" RC=0
  if command -v script > /dev/null 2>&1; then
    OUT="$(script -qec "$*" /dev/null 2>&1 </dev/null)" && RC=0 || RC=$?
  elif command -v python3 > /dev/null 2>&1; then
    OUT="$(python3 -c '
import os, pty, sys
st = pty.spawn(sys.argv[1:])
try: sys.exit(os.waitstatus_to_exitcode(st))
except AttributeError: sys.exit(st >> 8)' "$@" 2>&1 </dev/null)" && RC=0 || RC=$?
  else
    OUT="no pty tool available"
    RC=127
  fi
  return 0
}

# "nothing was started", as a hermetic question. Counting every sandbox container on the
# machine fails when the developer has a session open, or when an interrupted run left one
# behind -- neither of which is this test starting a container. Take the set of ids before,
# and ask whether it changed.
snapshot_containers() { AISB_CONTAINER_SNAPSHOT="$(podman ps -aq --filter label=ai-sandbox=1 2>/dev/null | sort)"; }
assert_no_new_containers() { # assert_no_new_containers <label>  (after snapshot_containers)
  local now extra
  now="$(podman ps -aq --filter label=ai-sandbox=1 2>/dev/null | sort)"
  if [[ "$now" == "${AISB_CONTAINER_SNAPSHOT:-}" ]]; then _ok "$1"; return 0; fi
  extra="$(comm -13 <(printf '%s\n' "${AISB_CONTAINER_SNAPSHOT:-}") \
                   <(printf '%s\n' "$now") | tr '\n' ' ')"
  _bad "$1" "container ids appeared: ${extra:-none (one disappeared instead)}"
}

# sb_env VAR=val ... -- args   (env prefix, without leaking into the test shell)
sb_env() {
  local -a envs=()
  while (($#)) && [[ "$1" != -* && "$1" == *=* ]]; do envs+=("$1"); shift; done
  OUT="$(env "${envs[@]}" "$AISB" "$@" 2>&1)" && RC=0 || RC=$?
  return 0
}

assert_rc() { # assert_rc <expected> <label>
  if [[ "$RC" == "$1" ]]; then _ok "$2"
  else _bad "$2" "expected rc=$1 got rc=$RC; output: $(printf '%s' "$OUT" | head -3 | tr '\n' ' ')"; fi
}

assert_contains() { # assert_contains <needle> <label>
  if printf '%s' "$OUT" | grep -qF -- "$1"; then _ok "$2"
  else _bad "$2" "missing: '$1'  in: $(printf '%s' "$OUT" | head -5 | tr '\n' ' ')"; fi
}

assert_absent() { # assert_absent <needle> <label>
  if printf '%s' "$OUT" | grep -qF -- "$1"; then
    _bad "$2" "unexpectedly present: '$1'"
  else _ok "$2"; fi
}

assert_eq() { # assert_eq <expected> <actual> <label>
  if [[ "$1" == "$2" ]]; then _ok "$3"
  else _bad "$3" "expected '$1' got '$2'"; fi
}

assert_count() { # assert_count <expected> <needle> <label>
  local n
  n="$(printf '%s' "$OUT" | grep -oF -- "$2" | wc -l | tr -d ' ')"
  if [[ "$n" == "$1" ]]; then _ok "$3"
  else _bad "$3" "expected $1 occurrence(s) of '$2', got $n"; fi
}

assert_file_mode() { # assert_file_mode <path> <mode> <label>
  local m; m="$(stat -c '%a' "$1" 2>/dev/null || echo missing)"
  if [[ "$m" == "$2" ]]; then _ok "$3"; else _bad "$3" "expected mode $2 got $m"; fi
}

summary() {
  printf '\n------------------------------------------------------------\n'
  [[ -n "${AISB_TEST_COUNTS:-}" ]] && printf '%s %s\n' "$PASS" "$FAIL" >> "$AISB_TEST_COUNTS"
  if (( FAIL )); then
    printf '\033[31m%d passed, %d FAILED\033[0m\n' "$PASS" "$FAIL"
    local f; for f in "${FAILED_NAMES[@]}"; do printf '  - %s\n' "$f"; done
    return 1
  fi
  printf '\033[32m%d passed, 0 failed\033[0m\n' "$PASS"
}