#!/usr/bin/env bash
# worktree_prepare.sh — prepare an issue worktree next to the control root
# (Issue #21 item 4). Run it from the control root's skills tree: the control
# root is lib.sh's WORKSPACE_ROOT, the main checkout that owns .mam.
#
# Usage:
#   worktree_prepare.sh --path <WT> --branch <name> [--base <ref>] [--submodules]
#                       [--creator-agent <agent> [--session <name>] [--role <role>]
#                        [--herdr-session <s>] [--herdr-workspace <w>] [--onboard]] [--dry-run]
#
# Steps: git worktree (idempotent) -> the worktree's own .venv (control root
# interpreter + requirements*.txt) -> /.venv and .pytest_cache/ in the shared info/exclude -> submodules only with
# --submodules (a failed init is rolled back: it would otherwise leave a
# modules/ trace that makes a plain `git worktree remove` refuse) -> optional
# creator session registered in the control root's .mam -> summary with BASE=.
#
# Exit codes: 0 ok, 1 git / submodule / create failure, 2 usage or path rule,
# 3 refused to protect existing content (nothing changed); 44/45 are create's
# (onboarding error / timeout; session and worktree kept, Issue #50).
set -euo pipefail

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

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

WT=""; BRANCH=""; BASE_REF=""; SUBMODULES=0; DRY_RUN=0
CREATOR_AGENT=""; SESSION=""; ROLE="creator"; HS=""; HW=""; ONBOARD=0
while [ $# -gt 0 ]; do
  case "$1" in
    --path) WT="${2:-}"; shift 2 ;;
    --branch) BRANCH="${2:-}"; shift 2 ;;
    --base) BASE_REF="${2:-}"; shift 2 ;;
    --submodules) SUBMODULES=1; shift ;;
    --creator-agent) CREATOR_AGENT="${2:-}"; shift 2 ;;
    --session) SESSION="${2:-}"; shift 2 ;;
    --role) ROLE="${2:-}"; shift 2 ;;
    --herdr-session) HS="${2:-}"; shift 2 ;;
    --herdr-workspace) HW="${2:-}"; shift 2 ;;
    --onboard) ONBOARD=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" ] && [ -n "$BRANCH" ] || { echo "ERROR: --path and --branch are required" >&2; usage >&2; exit 2; }

MAIN="$(mam_abs_workspace "$WORKSPACE_ROOT")"
top="$(git -C "$MAIN" rev-parse --show-toplevel 2>/dev/null || true)"
if [ -z "$top" ] || [ "$(mam_abs_workspace "$top")" != "$MAIN" ]; then
  echo "ERROR: the control root '$MAIN' (WORKSPACE_ROOT) is not the top of a git checkout" >&2
  exit 2
fi

case "$WT" in /*) ;; *) WT="$PWD/$WT" ;; esac
WT="${WT%/}"
if [ -d "$WT" ]; then
  WT="$(mam_abs_workspace "$WT")"
elif [ -d "$(dirname "$WT")" ]; then
  WT="$(mam_abs_workspace "$(dirname "$WT")")/$(basename "$WT")"
fi
case "$WT/" in
  "$MAIN"/*)
    echo "ERROR: '$WT' is inside the control root '$MAIN'; a worktree there shows up as untracked files in the main tree. Put it next to the control root instead." >&2
    exit 2 ;;
esac

PY="$(_delegate_py_bin)"
wt_py() { PYTHONPATH="$SKILL_DIR:${PYTHONPATH:-}" "$PY" -m lib_py.worktree "$@"; }

if [ -z "$SESSION" ] && [ -n "$CREATOR_AGENT" ]; then
  slug="$(printf '%s' "$BRANCH" | tr '[:upper:]' '[:lower:]' | tr -c 'a-z0-9-' '-' | sed -E 's/-+/-/g; s/^-//; s/-$//')"
  SESSION="creator-${CREATOR_AGENT}-${slug}"
  SESSION="${SESSION:0:60}"
  SESSION="${SESSION%-}"
fi

if [ "$DRY_RUN" = "1" ]; then
  echo "[dry-run] control root: $MAIN"
  echo "[dry-run] worktree: $WT (branch $BRANCH${BASE_REF:+, base $BASE_REF})"
  echo "[dry-run] own .venv (from $MAIN/.venv/bin/python + the delegate-job requirements.txt and requirements*.txt) and shared info/exclude entries; submodules: $([ "$SUBMODULES" = 1 ] && echo init || echo skipped)"
  [ -n "$CREATOR_AGENT" ] && echo "[dry-run] creator: $CREATOR_AGENT as $SESSION (role $ROLE)"
  exit 0
fi

# 1. The worktree (idempotent).
if [ "$(wt_py is-linked "$MAIN" "$WT")" = "true" ]; then
  cur="$(git -C "$WT" branch --show-current)"
  if [ "$cur" != "$BRANCH" ]; then
    echo "ERROR: '$WT' is already a worktree, on branch '${cur:-(detached)}', not '$BRANCH'" >&2
    exit 2
  fi
  echo "worktree already prepared: $WT ($BRANCH)"
else
  if [ -e "$WT" ] && [ -n "$(ls -A "$WT" 2>/dev/null)" ]; then
    echo "ERROR: '$WT' exists and is not a worktree of '$MAIN'" >&2
    exit 2
  fi
  mkdir -p "$(dirname "$WT")"
  if git -C "$MAIN" show-ref --verify --quiet "refs/heads/$BRANCH"; then
    if [ -n "$BASE_REF" ]; then
      echo "WARNING: branch '$BRANCH' already exists; --base '$BASE_REF' is ignored (the worktree checks the branch out as it is)" >&2
    fi
    git -C "$MAIN" worktree add "$WT" "$BRANCH" || exit 1
  else
    git -C "$MAIN" worktree add -b "$BRANCH" "$WT" "${BASE_REF:-HEAD}" || exit 1
  fi
fi

# 2. .venv: the worktree's own environment, created with the control root's
#    interpreter and its requirements files, never a link (review R3: a link
#    aliases the control root's environment, and a `git add -A` in the
#    worktree once committed the link itself - see .gitignore's managed
#    block). Idempotent: an existing usable venv is kept.
VENV_STATE="skipped"
if [ -L "$WT/.venv" ]; then
  echo "ERROR: '$WT/.venv' is a symbolic link; a worktree needs its own venv - remove the link and re-run" >&2
  exit 1
elif [ -x "$WT/.venv/bin/python" ]; then
  VENV_STATE="kept"
elif [ -x "$MAIN/.venv/bin/python" ]; then
  "$MAIN/.venv/bin/python" -m venv "$WT/.venv" || { echo "ERROR: could not create '$WT/.venv'" >&2; exit 1; }
  req=()
  # the same files deploy/install.sh installs into the control root's venv
  for f in .agents/skills/multi-agent-mux-delegate-job/requirements.txt requirements.txt requirements-dev.txt; do
    [ -f "$WT/$f" ] && req+=(-r "$WT/$f")
  done
  if [ "${#req[@]}" -gt 0 ]; then
    if ! "$WT/.venv/bin/python" -m pip install -q "${req[@]}"; then
      rm -rf "$WT/.venv"  # review D1: a half-made venv would be reported as kept by the next run
      echo "ERROR: pip install into '$WT/.venv' failed (network to the package index is needed); the venv was removed, the worktree is kept - fix it and re-run" >&2
      exit 1
    fi
  fi
  VENV_STATE="created"
else
  echo "WARNING: '$MAIN/.venv/bin/python' not found; no venv created in the worktree" >&2
fi

# 3. Ignore rules. A linked worktree's info/exclude is the shared
#    $GIT_COMMON_DIR/info/exclude, so these lines apply to the main checkout too
#    (they name artefacts it ignores as well).
excl="$(cd "$WT" && git rev-parse --path-format=absolute --git-path info/exclude)"
mkdir -p "$(dirname "$excl")"
touch "$excl"
if [ -s "$excl" ] && [ -n "$(tail -c1 "$excl")" ]; then echo >> "$excl"; fi
for pat in "/.venv" ".pytest_cache/"; do
  grep -Fxq -- "$pat" "$excl" || printf '%s\n' "$pat" >> "$excl"
done

# 4. Submodules: opt-in only.
SUB_STATE="skipped"
if [ -f "$WT/.gitmodules" ]; then
  if [ "$SUBMODULES" = "1" ]; then
    gd="$(git -C "$WT" rev-parse --absolute-git-dir)"
    had_modules=0
    if [ -e "$gd/modules" ]; then had_modules=1; fi
    # The state before this attempt, so a failure undoes only what it created
    # (review R3): per submodule its path and whether its git dir already
    # existed (an earlier init - kept), and the exact shared config values
    # (init and deinit both write the common .git/config). Indexed arrays only
    # (bash 3.2 has no associative ones).
    sm_names=(); sm_paths=(); sm_pre=(); cfg_keys=(); cfg_had=(); cfg_vals=()
    while read -r key path; do
      [ -n "$key" ] || continue
      name="${key#submodule.}"; name="${name%.path}"
      sm_names+=("$name"); sm_paths+=("$path")
      if [ -e "$gd/modules/$name" ]; then sm_pre+=(1); else sm_pre+=(0); fi
      for k in url active; do
        cfg_keys+=("submodule.$name.$k")
        if v="$(git -C "$MAIN" config --get "submodule.$name.$k" 2>/dev/null)"; then
          cfg_had+=(1); cfg_vals+=("$v")
        else
          cfg_had+=(0); cfg_vals+=("")
        fi
      done
    done < <(git config --file "$WT/.gitmodules" --get-regexp '^submodule\..*\.path$' 2>/dev/null || true)
    # A submodule path that is not initialised yet but already holds files
    # (review R4): git refuses to clone into it, and a rollback could not tell
    # those files from its own - refuse before changing anything. With this,
    # everything under a not-yet-initialised path after a failed attempt was
    # created by that attempt, which is what the rollback relies on.
    blocked=()
    for ((i = 0; i < ${#sm_names[@]}; i++)); do
      p="$WT/${sm_paths[$i]}"
      if [ "${sm_pre[$i]}" != "1" ] && [ -e "$p" ] && [ -n "$(ls -A "$p" 2>/dev/null)" ]; then
        blocked+=("${sm_paths[$i]}")
      fi
    done
    if [ "${#blocked[@]}" -gt 0 ]; then
      echo "REFUSED: submodule path(s) not initialised yet but not empty: ${blocked[*]}. Initialising would fail on them and a rollback could not tell their files from its own; move the files away (or drop --submodules). Nothing was changed." >&2
      exit 3
    fi
    if ! git -C "$WT" submodule update --init --recursive; then
      # A submodule cloned by this attempt leaves a checkout whose .git file
      # points into modules/<name>: deinit it before its git dir goes, or the
      # worktree stops being a usable git checkout.
      for ((i = 0; i < ${#sm_names[@]}; i++)); do
        if [ "${sm_pre[$i]}" = "1" ]; then continue; fi
        git -C "$WT" submodule deinit -q -f -- "${sm_paths[$i]}" >/dev/null 2>&1 || true
        if [ -f "$WT/${sm_paths[$i]}/.git" ]; then rm -f "$WT/${sm_paths[$i]}/.git"; fi
        rm -rf "$gd/modules/${sm_names[$i]}"
      done
      if [ "$had_modules" != "1" ]; then rm -rf "$gd/modules"; fi
      for ((i = 0; i < ${#cfg_keys[@]}; i++)); do
        if [ "${cfg_had[$i]}" = "1" ]; then
          git -C "$MAIN" config "${cfg_keys[$i]}" "${cfg_vals[$i]}"
        else
          git -C "$MAIN" config --unset-all "${cfg_keys[$i]}" 2>/dev/null || true
        fi
      done
      if st_err="$(git -C "$WT" status --porcelain 2>&1 >/dev/null)"; then
        echo "ERROR: submodule initialisation failed; this attempt's traces (new submodule checkouts, their git dirs, new submodule.* config keys) were removed and earlier submodule state was kept, so the worktree stays usable and removable. Re-run without --submodules to finish preparing." >&2
      else
        echo "ERROR: submodule initialisation failed, and '$WT' is not a usable git checkout after the rollback: $st_err" >&2
        echo "  inspect it with: git -C '$WT' submodule status - do not force-remove it before checking its submodules" >&2
      fi
      exit 1
    fi
    SUB_STATE="initialised"
  else
    echo "NOTE: submodules were not initialised (pass --submodules to initialise them)"
  fi
fi

# 5. Optional creator, registered in the control root's .mam.
if [ -n "$CREATOR_AGENT" ]; then
  args=(--workspace "$WT" --agent "$CREATOR_AGENT" --role "$ROLE" --session "$SESSION")
  if [ "$ONBOARD" = "1" ]; then args+=(--onboard); else args+=(--no-onboard); fi
  if [ -n "$HS" ]; then args+=(--herdr-session "$HS"); fi
  if [ -n "$HW" ]; then args+=(--herdr-workspace "$HW"); fi
  env_add=(WORKSPACE_ROOT="$MAIN"
           AGENT_SESSIONS_YAML="${AGENT_SESSIONS_YAML:-$MAIN/.mam/agent-sessions.yaml}"
           MAM_JOBS_DIR="${MAM_JOBS_DIR:-$MAIN/.mam/jobs}"
           MAM_NO_SUBSCRIBER_DAEMON=1)
  # Issue #57: an explicit audit-log root is passed on; otherwise create's
  # binding derives it from MAM_JOBS_DIR (the delegate_job_logs next to it).
  if [ -n "${DELEGATE_JOB_LOGS_DIR:-}" ]; then env_add+=(DELEGATE_JOB_LOGS_DIR="$DELEGATE_JOB_LOGS_DIR"); fi
  if [ -z "${MAM_ENV_FILE:-}" ] && [ -f "$MAIN/.mam.env" ]; then env_add+=(MAM_ENV_FILE="$MAIN/.mam.env"); fi
  rc=0
  env "${env_add[@]}" bash "$SKILL_DIR/multi-agent-mux-create/scripts/create_session.sh" "${args[@]}" || rc=$?
  if [ "$rc" -ne 0 ]; then
    case "$rc" in
      44) echo "ERROR: onboarding error for '$SESSION' (create_session.sh rc=44); session and worktree kept" >&2 ;;
      45) echo "ERROR: onboarding timeout for '$SESSION' (create_session.sh rc=45); session and worktree kept" >&2 ;;
      *) echo "ERROR: create_session.sh failed (rc=$rc); the worktree is kept - remove it with worktree_remove.sh --path '$WT'" >&2 ;;
    esac
    exit "$rc"
  fi
fi

BASE_SHA="$(git -C "$WT" rev-parse HEAD)"
echo "WORKTREE=$WT"
echo "BRANCH=$BRANCH"
echo "BASE=$BASE_SHA"
echo "SUBMODULES=$SUB_STATE"
echo "VENV=$VENV_STATE"
[ -n "$CREATOR_AGENT" ] && echo "SESSION=$SESSION"
echo "next: bash $SKILL_DIR/multi-agent-mux-loop/scripts/run_loop.sh --worktree '$WT' --base $BASE_SHA --review-only --creator <creator> --reviewer <reviewers> --task '...'"
echo "then: bash $SKILL_DIR/multi-agent-mux-stop/scripts/worktree_remove.sh --path '$WT' --stop"