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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MAIN_BRANCH="${SWARMFORGE_SCRIPTS_BRANCH:-main}"
ROLES_BRANCH="${SWARMFORGE_ROLES_BRANCH:-$MAIN_BRANCH}"
_REPO_BASE="${SWARMFORGE_REPO_URL:-https://github.com/gabadi/swarm-forge}"
ARCHIVE_URL="${SWARMFORGE_SCRIPTS_URL:-${_REPO_BASE}/archive/refs/heads/${MAIN_BRANCH}.tar.gz}"
ROLES_ARCHIVE_URL="${SWARMFORGE_ROLES_URL:-${_REPO_BASE}/archive/refs/heads/${ROLES_BRANCH}.tar.gz}"
TMP_DIR=""
ROLES_TMP_DIR=""

cleanup() {
  [[ -n "$TMP_DIR" ]] && rm -rf "$TMP_DIR"
  [[ -n "$ROLES_TMP_DIR" ]] && rm -rf "$ROLES_TMP_DIR"
}
trap cleanup EXIT

if [[ "${1:-}" == "stop" ]]; then
  STATE_DIR="$SCRIPT_DIR/.swarmforge"
  CLEANUP_SCRIPT="$SCRIPT_DIR/swarmforge/scripts/swarm-cleanup.sh"
  if [[ ! -f "$STATE_DIR/tmux-socket" ]]; then
    echo "No running swarm found (.swarmforge/tmux-socket missing)." >&2
    exit 1
  fi
  if [[ ! -f "$CLEANUP_SCRIPT" ]]; then
    echo "swarm-cleanup.sh not found — run ./swarm once to bootstrap scripts." >&2
    exit 1
  fi
  _tmux_socket="$(< "$STATE_DIR/tmux-socket")"
  _sessions=()
  if [[ -f "$STATE_DIR/sessions.tsv" ]]; then
    while IFS=$'\t' read -r _ _ _session _; do
      [[ -n "${_session:-}" ]] && _sessions+=("$_session")
    done < "$STATE_DIR/sessions.tsv"
  fi
  SWARMFORGE_TERMINAL_BACKEND="${SWARMFORGE_TERMINAL_BACKEND:-${SWARMFORGE_TERMINAL:-ghostty}}" \
    "$CLEANUP_SCRIPT" \
    "$_tmux_socket" \
    "$STATE_DIR/window-ids" \
    "${_sessions[@]+"${_sessions[@]}"}"
  echo "SwarmForge stopped."
  exit 0
fi

if [[ "${1:-}" == "upgrade" ]]; then
  echo "Upgrading swarmforge scripts from $MAIN_BRANCH, roles from $ROLES_BRANCH..."

  # Re-fetch scripts (force by removing existing)
  rm -rf "$SCRIPT_DIR/swarmforge/scripts"
  TMP_DIR="$(mktemp -d)"
  curl -L "$ARCHIVE_URL" | tar -xz --strip-components=1 -C "$TMP_DIR"
  cp -R "$TMP_DIR/swarmforge/scripts" "$SCRIPT_DIR/swarmforge/scripts"

  # Overwrite shared articles from the downloaded archive
  if [[ -d "$TMP_DIR/swarmforge/constitution/articles" ]]; then
    mkdir -p "$SCRIPT_DIR/swarmforge/scripts/shared-articles"
    cp -R "$TMP_DIR/swarmforge/constitution/articles/." "$SCRIPT_DIR/swarmforge/scripts/shared-articles/"
    mkdir -p "$SCRIPT_DIR/swarmforge/constitution/articles"
    for _article in "$SCRIPT_DIR/swarmforge/scripts/shared-articles"/*.prompt; do
      cp "$_article" "$SCRIPT_DIR/swarmforge/constitution/articles/$(basename "$_article")"
    done
  fi

  # Update roles — use a separate archive if ROLES_BRANCH differs from MAIN_BRANCH
  if [[ "$ROLES_BRANCH" == "$MAIN_BRANCH" ]]; then
    _roles_src="$TMP_DIR"
  else
    ROLES_TMP_DIR="$(mktemp -d)"
    curl -L "$ROLES_ARCHIVE_URL" | tar -xz --strip-components=1 -C "$ROLES_TMP_DIR"
    _roles_src="$ROLES_TMP_DIR"
  fi
  if [[ -d "$_roles_src/swarmforge/roles" ]]; then
    cp -R "$_roles_src/swarmforge/roles/." "$SCRIPT_DIR/swarmforge/roles/"
  fi

  # Update skills
  if [[ -d "$_roles_src/swarmforge/skills" ]]; then
    rm -rf "$SCRIPT_DIR/swarmforge/skills"
    cp -R "$_roles_src/swarmforge/skills" "$SCRIPT_DIR/swarmforge/skills"
  fi

  echo "SwarmForge upgraded."
  exit 0
fi

if [[ ! -d "$SCRIPT_DIR/swarmforge/scripts" || ! -d "$SCRIPT_DIR/swarmforge/scripts/shared-articles" || ! -d "$SCRIPT_DIR/swarmforge/skills" ]]; then
  TMP_DIR="$(mktemp -d)"
  mkdir -p "$SCRIPT_DIR/swarmforge"
  curl -L "$ARCHIVE_URL" | tar -xz --strip-components=1 -C "$TMP_DIR"
  if [[ ! -d "$SCRIPT_DIR/swarmforge/scripts" ]]; then
    cp -R "$TMP_DIR/swarmforge/scripts" "$SCRIPT_DIR/swarmforge/scripts"
  fi
  if [[ -d "$TMP_DIR/swarmforge/constitution/articles" ]]; then
    mkdir -p "$SCRIPT_DIR/swarmforge/scripts/shared-articles"
    cp -R "$TMP_DIR/swarmforge/constitution/articles/." "$SCRIPT_DIR/swarmforge/scripts/shared-articles/"
  fi
  if [[ ! -d "$SCRIPT_DIR/swarmforge/skills" && -d "$TMP_DIR/swarmforge/skills" ]]; then
    cp -R "$TMP_DIR/swarmforge/skills" "$SCRIPT_DIR/swarmforge/skills"
  fi
fi

# Copy shared articles into constitution/articles/ if not already present.
# Runs unconditionally — projects installed from a pre-packaged branch already
# ship scripts/shared-articles/ so the download block above is skipped, but the
# articles still need to land in constitution/articles/ for agent bundles to include them.
if [[ -d "$SCRIPT_DIR/swarmforge/scripts/shared-articles" ]]; then
  mkdir -p "$SCRIPT_DIR/swarmforge/constitution/articles"
  for _article in "$SCRIPT_DIR/swarmforge/scripts/shared-articles"/*.prompt; do
    _dest="$SCRIPT_DIR/swarmforge/constitution/articles/$(basename "$_article")"
    [[ ! -f "$_dest" ]] && cp "$_article" "$_dest"
  done
fi

# Install local skills immediately so /setup-swarm is available in Claude Code
# before swarmforge.sh runs and before setup-complete exists.
if [[ -d "$SCRIPT_DIR/swarmforge/skills" ]]; then
  mkdir -p "$SCRIPT_DIR/.claude/skills"
  for _skill_dir in "$SCRIPT_DIR/swarmforge/skills"/*/; do
    _skill_name="$(basename "$_skill_dir")"
    cp -R "$_skill_dir" "$SCRIPT_DIR/.claude/skills/$_skill_name"
  done
fi

exec "$SCRIPT_DIR/swarmforge/scripts/swarmforge.sh" "$@"
