#!/usr/bin/env bash
# worktree_remove.sh — remove an issue worktree of the control root (Issue #21
# item 4). Every check runs first and changes nothing; only when all pass are
# sessions stopped and the worktree removed.
#
# Usage:
#   worktree_remove.sh --path <WT> [--stop] [--dry-run]
#
# Checks: (1) a linked worktree of the control root, not locked; (2) no running
# session under it unless --stop; (3) no open job of those sessions (the
# delegate job of a session --stop will stop is cancelled by stop itself);
# (4) clean tree, submodules included; (5) submodules: nothing modified and no
# commit that is on no remote (HEAD included - a submodule is on a detached
# HEAD) and no stash. (4) and (5) run again right before the removal, after
# --stop: a session that is being stopped can still write, commit or stash.
# Removal is plain, or - when the worktree holds submodule state, which
# git refuses to remove without it - a single --force justified by (4) and (5).
# Never run `git worktree remove --force` by hand on such a worktree.
#
# Exit codes: 0 removed, 1 stop/git failure, 2 not a worktree of the control
# root, 3 refused by a check (nothing changed).
set -euo pipefail

_script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$(cd "$_script_dir/../.." && pwd)/lib.sh"

usage() {
  sed -n '2,20p' "${BASH_SOURCE[0]}" | sed 's/^# \{0,1\}//'
}

WT=""; STOP=0; DRY_RUN=0
while [ $# -gt 0 ]; do
  case "$1" in
    --path) WT="${2:-}"; shift 2 ;;
    --stop) STOP=1; shift ;;
    --dry-run) DRY_RUN=1; shift ;;
    -h|--help) usage; exit 0 ;;
    *) echo "ERROR: unknown arg: $1" >&2; usage >&2; exit 2 ;;
  esac
done
[ -n "$WT" ] || { echo "ERROR: --path is required" >&2; usage >&2; exit 2; }

MAIN="$(mam_abs_workspace "$WORKSPACE_ROOT")"
case "$WT" in /*) ;; *) WT="$PWD/$WT" ;; esac
WT="$(mam_abs_workspace "${WT%/}")"

PY="$(_delegate_py_bin)"
wt_py() { PYTHONPATH="$SKILL_DIR:${PYTHONPATH:-}" "$PY" -m lib_py.worktree "$@"; }
jget() { python3 -c 'import json,sys; d=json.load(sys.stdin); print(eval(sys.argv[1]))' "$1"; }

if [ "$(wt_py is-linked "$MAIN" "$WT")" != "true" ]; then
  echo "ERROR: '$WT' is not a linked worktree of the control root '$MAIN'" >&2
  exit 2
fi

REFUSED=0
refuse() { echo "REFUSED: $*" >&2; REFUSED=1; }

# (1) locked: even a single --force is refused by git
if [ "$(wt_py is-locked "$MAIN" "$WT")" = "true" ]; then
  refuse "the worktree is locked (run: git worktree unlock '$WT')"
fi

# (2) sessions
STATE="$(load_state_json)"
ROWS="$(MAM_STATE_JSON="$STATE" wt_py rows "$WT")"
RUNNING="$(printf '%s' "$ROWS" | jget '" ".join(s["name"] for s in d if s.get("status") == "running")')"
RUNNING_AGENTS="$(printf '%s' "$ROWS" | jget '" ".join(s["name"] + ":" + ((s.get("pane") or {}).get("cmd") or "") for s in d if s.get("status") == "running")')"
STOP_JIDS="$(printf '%s' "$ROWS" | jget '",".join(str(s.get("delegate_job_id")) for s in d if s.get("status") == "running" and s.get("delegate_job_id"))')"
if [ -n "$RUNNING" ] && [ "$STOP" != "1" ]; then
  refuse "running session(s) in the worktree: $RUNNING (pass --stop to stop them before removal)"
fi

# (3) open jobs (not the delegate jobs --stop will cancel)
EXCLUDE=""
[ "$STOP" = "1" ] && EXCLUDE="$STOP_JIDS"
JOBS="$(MAM_STATE_JSON="$STATE" wt_py open-jobs "$WT" --exclude "$EXCLUDE")"
if [ "$(printf '%s' "$JOBS" | jget 'len(d)')" != "0" ]; then
  refuse "open job(s) of sessions in the worktree - not cancelled automatically, they may belong to a running loop:"
  printf '%s' "$JOBS" | jget '"\n".join("  %s  %s  %s  (%s)" % (j["job_id"], j["agent_session"], j["status"], j["registry_dir"]) for j in d)' >&2
  echo "  close one after checking it: python3 .agents/skills/multi-agent-mux-delegate-job/scripts/registry.py --registry-dir <dir> close --job <id> --request cancelled" >&2
fi

# (4) clean, submodules included, and (5) submodules. A function: it runs now,
# before any change, and again right before the removal (review R1). It sets
# METHOD from what it saw and returns 1 on any finding.
judge_tree() {
  local found=0 dirty sub trace key val
  if ! dirty="$(git -C "$WT" status --porcelain --ignore-submodules=none 2>&1)"; then
    refuse "git status failed in the worktree: $dirty"
    return 1
  fi
  if [ -n "$dirty" ]; then
    refuse "the worktree has uncommitted changes:"
    printf '%s\n' "$dirty" | sed 's/^/  /' >&2
    found=1
  fi
  sub="$(wt_py submodules "$WT")"
  trace="$(printf '%s' "$sub" | jget 'd["trace"]')"
  METHOD="plain"
  if [ "$trace" = "True" ]; then
    METHOD="single --force"
    for key in error dirty unpushed stashes; do
      val="$(printf '%s' "$sub" | jget "d.get('$key', '')")"
      if [ -n "$val" ]; then
        refuse "submodule check '$key' is not empty (a forced removal would destroy it):"
        printf '%s\n' "$val" | sed 's/^/  /' >&2
        found=1
      fi
    done
  fi
  return "$found"
}
judge_tree || true

echo "checks: locked/sessions/jobs/clean/submodules -> $([ "$REFUSED" = 1 ] && echo refused || echo passed); removal method: $METHOD"
if [ "$REFUSED" = "1" ]; then
  echo "nothing was changed" >&2
  exit 3
fi
if [ "$DRY_RUN" = "1" ]; then
  echo "[dry-run] would ${RUNNING:+stop: $RUNNING, then }remove '$WT' ($METHOD)"
  exit 0
fi

# (6) stop, then re-check every open job
if [ "$STOP" = "1" ] && [ -n "$RUNNING_AGENTS" ]; then
  for pair in $RUNNING_AGENTS; do
    name="${pair%%:*}"; agent="${pair#*:}"
    stop_args=(--session "$name")
    [ -n "$agent" ] && stop_args+=(--agent "$agent")
    if ! bash "$SKILL_DIR/multi-agent-mux-stop/scripts/stop_session.sh" "${stop_args[@]}"; then
      echo "ERROR: stopping '$name' failed; the worktree was not removed" >&2
      exit 1
    fi
  done
  JOBS="$(MAM_STATE_JSON="$(load_state_json)" wt_py open-jobs "$WT")"
  if [ "$(printf '%s' "$JOBS" | jget 'len(d)')" != "0" ]; then
    echo "REFUSED: the sessions were stopped, but open job(s) appeared meanwhile; the worktree was not removed:" >&2
    printf '%s' "$JOBS" | jget '"\n".join("  %s  %s  %s" % (j["job_id"], j["agent_session"], j["status"]) for j in d)' >&2
    exit 3
  fi
fi

# (6b) the tree again, right before removing it: the stopped sessions may have
#      written, committed or stashed while they stopped (review R1). The
#      removal method follows this latest look.
if ! judge_tree; then
  echo "REFUSED: ${RUNNING:+the sessions were stopped, but }the worktree changed after the checks; it was not removed" >&2
  exit 3
fi

# (7) remove; the branch is kept
if [ "$METHOD" = "plain" ]; then
  git -C "$MAIN" worktree remove "$WT" || exit 1
else
  git -C "$MAIN" worktree remove --force "$WT" || exit 1
  echo "submodule checks passed; removed with a single --force"
fi
echo "removed: $WT ($METHOD)"
NAMES="$(printf '%s' "$ROWS" | jget '" ".join(s["name"] for s in d)')"
if [ -n "$NAMES" ]; then
  echo "Sessions created in this worktree ($NAMES) cannot be resumed; their conversations are kept, and recreating a worktree at the same path makes them resumable again."
fi