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

SCRIPT_DIR=$(cd -- "$(dirname -- "$0")" && pwd)
START="$SCRIPT_DIR/start-agent-session.sh"
SOCKET="mosaic-agent-test-$RANDOM-$$"
AGENT="agent-$RANDOM"
WORKDIR=$(mktemp -d)

# Keep a single cleanup trap that accumulates resources.
CLEANUP_DIRS=("$WORKDIR")
CLEANUP_SOCKETS=("$SOCKET")
trap '_cleanup' EXIT
_cleanup() {
  for s in "${CLEANUP_SOCKETS[@]:-}"; do
    tmux -L "$s" kill-server >/dev/null 2>&1 || true
  done
  for d in "${CLEANUP_DIRS[@]:-}"; do
    rm -rf "$d"
  done
}

fail() {
  echo "FAIL: $*" >&2
  exit 1
}

# ── Test 1: basic session creation with workdir check ─────────────────────────
MOSAIC_TMUX_SOCKET="$SOCKET" \
MOSAIC_AGENT_WORKDIR="$WORKDIR" \
MOSAIC_AGENT_COMMAND='bash --noprofile --norc -i' \
  "$START" "$AGENT"

tmux -L "$SOCKET" has-session -t "=$AGENT:0.0" || fail "agent session was not created"
actual_dir=$(tmux -L "$SOCKET" display-message -p -t "=$AGENT:0.0" '#{pane_current_path}')
[ "$actual_dir" = "$WORKDIR" ] || fail "agent workdir mismatch: $actual_dir"

# ── Test 2: idempotency (duplicate start prints 'already running') ─────────────
MOSAIC_TMUX_SOCKET="$SOCKET" \
MOSAIC_AGENT_WORKDIR="$WORKDIR" \
MOSAIC_AGENT_COMMAND='bash --noprofile --norc -i' \
  "$START" "$AGENT" >/tmp/mosaic-start-agent-idempotent.out

grep -qF 'already running' /tmp/mosaic-start-agent-idempotent.out || fail "duplicate start was not idempotent"

# ── Test 3: runtime-bin PATH prefix is baked into the pane command ────────────
#
# We capture the command the script would hand to tmux by injecting a fake
# 'tmux' shim into PATH.  The shim:
#   - Intercepts 'new-session' calls and records its arguments to a file.
#   - For 'has-session' calls, exits 1 (session does not exist) so the script
#     proceeds to launch instead of printing "already running".
#   - For 'list-panes' calls, returns empty so PANE_PID stays unset and the
#     heartbeat sidecar is NOT spawned (heartbeat is not the focus of this test;
#     test 6 and 7 cover that path).  This prevents any real-filesystem side
#     effects or leaked background processes.
#   - For all other subcommands, exits 0.
#
# Assertions:
#   a) 'export PATH=' with the synthetic MOSAIC_RUNTIME_BIN prefix appears.
#   b) 'exec' appears so the runtime replaces the wrapper shell.
#   c) MOSAIC_AGENT_COMMAND with flags is forwarded intact.

FAKE_BIN=$(mktemp -d)
FAKE_RUNTIME_BIN=$(mktemp -d)
TMUX_ARGS_FILE=$(mktemp)
HB_RUN_DIR3=$(mktemp -d)
CLEANUP_DIRS+=("$FAKE_BIN" "$FAKE_RUNTIME_BIN" "$HB_RUN_DIR3")

# Write the fake tmux shim (uses only positional args, no sourced vars).
cat > "$FAKE_BIN/tmux" <<SHIM
#!/usr/bin/env bash
# Fake tmux: record new-session args; report has-session as missing.
subcmd="\$3"  # argv: tmux -L <socket> <subcmd> ...
if [ "\$subcmd" = "has-session" ]; then
  exit 1   # session not found → script will attempt new-session
fi
if [ "\$subcmd" = "new-session" ]; then
  printf '%s\n' "\$@" > "$TMUX_ARGS_FILE"
  exit 0
fi
if [ "\$subcmd" = "list-panes" ]; then
  # Return empty: no sidecar spawned (heartbeat is not the focus of this test).
  echo ""
  exit 0
fi
exit 0
SHIM
chmod +x "$FAKE_BIN/tmux"

SOCKET3="mosaic-agent-test3-$RANDOM-$$"
AGENT3="agent3-$RANDOM"
WORKDIR3=$(mktemp -d)
CLEANUP_DIRS+=("$WORKDIR3")

PATH="$FAKE_BIN:$PATH" \
MOSAIC_TMUX_SOCKET="$SOCKET3" \
MOSAIC_AGENT_WORKDIR="$WORKDIR3" \
MOSAIC_AGENT_RUNTIME="pi" \
MOSAIC_RUNTIME_BIN="$FAKE_RUNTIME_BIN" \
MOSAIC_AGENT_COMMAND="mosaic yolo pi --model openai-codex/gpt-5.5:high" \
MOSAIC_HEARTBEAT_RUN_DIR="$HB_RUN_DIR3" \
  "$START" "$AGENT3"

all_args=$(cat "$TMUX_ARGS_FILE" 2>/dev/null || true)
rm -f "$TMUX_ARGS_FILE"

echo "--- captured tmux new-session args ---"
echo "$all_args"
echo "--- end args ---"

# a) PATH prefix containing FAKE_RUNTIME_BIN must appear.
echo "$all_args" | grep -qF "export PATH=" || fail "pane command does not export PATH"
echo "$all_args" | grep -qF "$FAKE_RUNTIME_BIN" || fail "pane command does not include MOSAIC_RUNTIME_BIN in PATH prefix"

# b) exec must appear so the runtime replaces the wrapper shell.
echo "$all_args" | grep -qF "exec " || fail "pane command does not use exec"

# c) Full MOSAIC_AGENT_COMMAND (with flags) must be forwarded.
echo "$all_args" | grep -qF "mosaic yolo pi --model openai-codex/gpt-5.5:high" || \
  fail "pane command does not forward MOSAIC_AGENT_COMMAND with flags intact"

# ── Test 4: when no extra runtime-bin dirs exist, exec still appears ───────────
TMUX_ARGS_FILE2=$(mktemp)
FAKE_BIN2=$(mktemp -d)
HB_RUN_DIR4=$(mktemp -d)
CLEANUP_DIRS+=("$FAKE_BIN2" "$HB_RUN_DIR4")

cat > "$FAKE_BIN2/tmux" <<SHIM2
#!/usr/bin/env bash
subcmd="\$3"
if [ "\$subcmd" = "has-session" ]; then exit 1; fi
if [ "\$subcmd" = "new-session" ]; then
  printf '%s\n' "\$@" > "$TMUX_ARGS_FILE2"
  exit 0
fi
if [ "\$subcmd" = "list-panes" ]; then
  # Return empty: no sidecar spawned (heartbeat is not the focus of this test).
  echo ""
  exit 0
fi
exit 0
SHIM2
chmod +x "$FAKE_BIN2/tmux"

SOCKET4="mosaic-agent-test4-$RANDOM-$$"
AGENT4="agent4-$RANDOM"
WORKDIR4=$(mktemp -d)
CLEANUP_DIRS+=("$WORKDIR4")

# MOSAIC_RUNTIME_BIN points to a non-existent dir so prefix will be empty;
# .npm-global/bin and .local/bin may or may not exist but we just want exec.
PATH="$FAKE_BIN2:$PATH" \
MOSAIC_TMUX_SOCKET="$SOCKET4" \
MOSAIC_AGENT_WORKDIR="$WORKDIR4" \
MOSAIC_AGENT_RUNTIME="pi" \
MOSAIC_RUNTIME_BIN="/nonexistent-dir-$$" \
MOSAIC_AGENT_COMMAND="mosaic yolo pi" \
MOSAIC_HEARTBEAT_RUN_DIR="$HB_RUN_DIR4" \
  "$START" "$AGENT4"

all_args4=$(cat "$TMUX_ARGS_FILE2" 2>/dev/null || true)
rm -f "$TMUX_ARGS_FILE2"
rm -rf "$WORKDIR4"

echo "$all_args4" | grep -qF "exec " || fail "pane command (no prefix dirs) does not use exec"
echo "$all_args4" | grep -qF "mosaic yolo pi" || fail "pane command does not include agent command when no prefix"

# ── Test 5: candidate dir already in LAUNCHER $PATH is still baked into pane ──
#
# Regression guard for the bug where _build_runtime_bin_prefix() used to skip
# a candidate because it was already present in the launcher process's $PATH.
# That check was wrong: the pane inherits the tmux SERVER environment, not the
# launcher's env.  Even if a dir is on the launcher's PATH it must always be
# baked into the pane's PATH export.
#
# We prove this by setting PATH to include FAKE_RUNTIME_BIN5 (the candidate),
# then asserting the generated new-session command still exports it.
TMUX_ARGS_FILE5=$(mktemp)
FAKE_BIN5=$(mktemp -d)
FAKE_RUNTIME_BIN5=$(mktemp -d)  # this dir IS on the launcher's PATH below
HB_RUN_DIR5=$(mktemp -d)
CLEANUP_DIRS+=("$FAKE_BIN5" "$FAKE_RUNTIME_BIN5" "$HB_RUN_DIR5")

cat > "$FAKE_BIN5/tmux" <<SHIM5
#!/usr/bin/env bash
subcmd="\$3"
if [ "\$subcmd" = "has-session" ]; then exit 1; fi
if [ "\$subcmd" = "new-session" ]; then
  printf '%s\n' "\$@" > "$TMUX_ARGS_FILE5"
  exit 0
fi
if [ "\$subcmd" = "list-panes" ]; then
  # Return empty: no sidecar spawned (heartbeat is not the focus of this test).
  echo ""
  exit 0
fi
exit 0
SHIM5
chmod +x "$FAKE_BIN5/tmux"

SOCKET5="mosaic-agent-test5-$RANDOM-$$"
AGENT5="agent5-$RANDOM"
WORKDIR5=$(mktemp -d)
CLEANUP_DIRS+=("$WORKDIR5")
CLEANUP_SOCKETS+=("$SOCKET5")

# FAKE_RUNTIME_BIN5 is deliberately placed on the LAUNCHER PATH so that the
# old (buggy) code would have skipped it.  The correct code must still include
# it in the pane PATH export.
PATH="$FAKE_BIN5:$FAKE_RUNTIME_BIN5:$PATH" \
MOSAIC_TMUX_SOCKET="$SOCKET5" \
MOSAIC_AGENT_WORKDIR="$WORKDIR5" \
MOSAIC_AGENT_RUNTIME="pi" \
MOSAIC_RUNTIME_BIN="$FAKE_RUNTIME_BIN5" \
MOSAIC_AGENT_COMMAND="mosaic yolo pi" \
MOSAIC_HEARTBEAT_RUN_DIR="$HB_RUN_DIR5" \
  "$START" "$AGENT5"

all_args5=$(cat "$TMUX_ARGS_FILE5" 2>/dev/null || true)
rm -f "$TMUX_ARGS_FILE5"
rm -rf "$WORKDIR5"

echo "--- test 5: launcher-PATH candidate must still appear in pane export ---"
echo "$all_args5"
echo "--- end test 5 args ---"

echo "$all_args5" | grep -qF "export PATH=" || \
  fail "test5: pane command does not export PATH when candidate is on launcher PATH"
echo "$all_args5" | grep -qF "$FAKE_RUNTIME_BIN5" || \
  fail "test5: candidate dir (already on launcher PATH) was NOT baked into pane PATH — regression"

# ── Test 6: heartbeat sidecar — pane PID resolved + .hb file written ──────────
#
# Uses a real tmux session (same socket as test 1 which already has $AGENT) so
# list-panes returns a real pane PID.  We override MOSAIC_HEARTBEAT_RUN_DIR to
# a temp dir and set a 1-second interval, then wait up to 3 s for the .hb file
# to appear and check its content.

HB_RUN_DIR=$(mktemp -d)
CLEANUP_DIRS+=("$HB_RUN_DIR")

# Re-use the session+agent created in Test 1 (still alive on $SOCKET / $AGENT).
# We need to invoke the script for a NEW agent on the same socket to exercise
# the heartbeat path with a real pane PID.
AGENT6="agent6-$RANDOM"
MOSAIC_TMUX_SOCKET="$SOCKET" \
MOSAIC_AGENT_WORKDIR="$WORKDIR" \
MOSAIC_AGENT_COMMAND='bash --noprofile --norc -i' \
MOSAIC_HEARTBEAT_RUN_DIR="$HB_RUN_DIR" \
MOSAIC_HEARTBEAT_INTERVAL="1" \
  "$START" "$AGENT6"

HB_FILE="$HB_RUN_DIR/${AGENT6}.hb"

# Wait up to 5 seconds for the heartbeat file to appear.
_waited=0
until [ -f "$HB_FILE" ] || [ "$_waited" -ge 5 ]; do
  sleep 0.5
  _waited=$((_waited + 1))
done

[ -f "$HB_FILE" ] || fail "test6: heartbeat file not written at $HB_FILE within 5s"

hb_content=$(cat "$HB_FILE")
echo "--- test 6: heartbeat file content ---"
echo "$hb_content"
echo "--- end test 6 ---"

# Verify required fields are present.
echo "$hb_content" | grep -qE '^ts=[0-9]{4}-[0-9]{2}-[0-9]{2}T' || \
  fail "test6: heartbeat ts field missing or malformed"
echo "$hb_content" | grep -qE '^pid=[0-9]+' || \
  fail "test6: heartbeat pid field missing or malformed"
echo "$hb_content" | grep -qF 'status=ok' || \
  fail "test6: heartbeat status=ok missing"

# ── Test 7: heartbeat sidecar — targets correct .hb path per agent name ────────
#
# Uses the fake-tmux shim approach (like tests 3-5) to capture the sidecar
# invocation without needing a real session.  A fake setsid shim records its
# arguments so we can assert the sidecar script targets the expected .hb path
# and uses the configured interval.

FAKE_BIN7=$(mktemp -d)
FAKE_RUNTIME_BIN7=$(mktemp -d)
SETSID_ARGS_FILE=$(mktemp)
HB_RUN_DIR7=$(mktemp -d)
CLEANUP_DIRS+=("$FAKE_BIN7" "$FAKE_RUNTIME_BIN7" "$HB_RUN_DIR7")

AGENT7="my-fleet-agent-$RANDOM"
INTERVAL7="42"

# Fake tmux: has-session → not found; new-session → ok; list-panes → known PID.
cat > "$FAKE_BIN7/tmux" <<SHIM7
#!/usr/bin/env bash
subcmd="\$3"
if [ "\$subcmd" = "has-session" ]; then exit 1; fi
if [ "\$subcmd" = "new-session" ]; then exit 0; fi
if [ "\$subcmd" = "list-panes" ]; then echo "88888"; exit 0; fi
exit 0
SHIM7
chmod +x "$FAKE_BIN7/tmux"

# Fake setsid: capture the bash -c <script> argument for inspection, then
# background an actual bash subshell so disown succeeds in the caller.
cat > "$FAKE_BIN7/setsid" <<'SETSID_SHIM'
#!/usr/bin/env bash
# argv: setsid bash -c <sidecar_script>
# Record the full argument list to the capture file, then exit cleanly.
printf '%s\0' "$@" > __SETSID_ARGS_FILE__
exit 0
SETSID_SHIM
# Patch the placeholder with the real capture-file path (avoids heredoc expansion issues).
sed -i "s|__SETSID_ARGS_FILE__|${SETSID_ARGS_FILE}|g" "$FAKE_BIN7/setsid"
chmod +x "$FAKE_BIN7/setsid"

SOCKET7="mosaic-agent-test7-$RANDOM-$$"
WORKDIR7=$(mktemp -d)
CLEANUP_DIRS+=("$WORKDIR7")

PATH="$FAKE_BIN7:$PATH" \
MOSAIC_TMUX_SOCKET="$SOCKET7" \
MOSAIC_AGENT_WORKDIR="$WORKDIR7" \
MOSAIC_AGENT_RUNTIME="pi" \
MOSAIC_RUNTIME_BIN="$FAKE_RUNTIME_BIN7" \
MOSAIC_AGENT_COMMAND="mosaic yolo pi" \
MOSAIC_HEARTBEAT_RUN_DIR="$HB_RUN_DIR7" \
MOSAIC_HEARTBEAT_INTERVAL="$INTERVAL7" \
  "$START" "$AGENT7"

# Give the background setsid shim a moment to finish writing the capture file.
sleep 0.5

setsid_args=$(cat "$SETSID_ARGS_FILE" 2>/dev/null | tr '\0' '\n' || true)
rm -f "$SETSID_ARGS_FILE"
rm -rf "$WORKDIR7"

echo "--- test 7: captured setsid args ---"
echo "$setsid_args"
echo "--- end test 7 ---"

# The sidecar script (bash -c <script>) must reference the correct .hb path.
expected_hb="${HB_RUN_DIR7}/${AGENT7}.hb"
echo "$setsid_args" | grep -qF "$expected_hb" || \
  fail "test7: sidecar script does not reference correct .hb path ($expected_hb)"

# The sidecar script must use the configured interval.
echo "$setsid_args" | grep -qF "$INTERVAL7" || \
  fail "test7: sidecar script does not reference configured interval ($INTERVAL7)"

echo "ok - start-agent-session"