#!/bin/bash
# ai — agent session launcher (unified aix + hey)
#
# Usage:
#   ai                 Select one or more projects via fzf and open each in
#                      an 'ai' tmux session window (creates session if needed).
#                      Git-aware: fetches, annotates with ↑/↓/dirty, auto-pulls
#                      clean-and-behind repos before opening.
#
#   ai <dir>...        Single-project mode. Opens each given directory directly
#                      in the 'ai' session (new window or switch to existing).
#                      Use '.' for current directory. Git prep per dir.
#
#   ai --runtime <rt>  Launch with a different agent CLI: claude (default) or
#                      codex. 'local' is reserved for a local-LLM CLI, pending
#                      the model evaluation. Also settable via AI_RUNTIME.
#
#   ai --attach        Attach to the existing 'ai' session without changes.
#
#   ai -h | --help     Show this help.
#
# Source: ~/code/rulesets/claude-templates/bin/ai
# Install: make -C ~/code/rulesets install
#
# NOTE: do not enable `set -e`. Several helper functions use
# `[ test ] && action` patterns that legitimately return non-zero when
# the test fails (e.g. maybe_add_candidate when a dir lacks .ai/protocols.org,
# git_status_indicator when a repo is clean). With set -e those exit codes
# would kill the script.

SESSION="ai"
RUNTIME="${AI_RUNTIME:-claude}"

# Map the runtime name to the agent CLI a pane launches. Both current CLIs
# take the opening instructions as a positional prompt, so only the command
# name varies; a runtime whose CLI needs different plumbing gets its own case.
resolve_agent_cmd() {
  case "$RUNTIME" in
    claude) AGENT_CMD="claude" ;;
    codex)  AGENT_CMD="codex" ;;
    local)
      echo "ai: the 'local' runtime is reserved but not wired yet (pending the local-model evaluation) — valid runtimes: claude, codex" >&2
      exit 2
      ;;
    *)
      echo "ai: unknown runtime '$RUNTIME' — valid runtimes: claude, codex" >&2
      exit 2
      ;;
  esac
}

# Run in the pane's shell just before Claude launches. `stty susp undef` clears
# the tty's SIGTSTP (C-z) character for this pane only, so an accidental C-z is
# passed through to Claude as input rather than suspending the session to the
# shell. Scoped here so C-z keeps working as job control in every other
# terminal, shell, and program.
LAUNCH_PREFIX="stty susp undef; "

# Format the per-project opening line passed to claude. Takes the project
# directory's basename; returns a string of the form
#   "This is <host> <name> project. Follow all instructions in .ai/protocols.org."
# Uses uname -n for the host (POSIX, no dependency on the hostname binary
# which isn't installed by default on every distro). The project name
# disambiguates windows when multiple projects share an ai-session.
build_instructions() {
  local name="$1"
  printf 'This is %s %s project. Follow all instructions in .ai/protocols.org.' "$(uname -n)" "$name"
}

usage() {
  sed -n '2,20p' "$0" | sed 's|^# \?||'
  exit 0
}

check_deps() {
  for cmd in fzf tmux "$AGENT_CMD"; do
    if ! command -v "$cmd" &>/dev/null; then
      echo "ai: $cmd is not installed" >&2
      exit 1
    fi
  done
}

# ---------- shared helpers ----------

attach_session() {
  if [ -n "${TMUX:-}" ]; then
    tmux switch-client -t "$SESSION"
  else
    tmux attach-session -t "$SESSION"
  fi
}

# Create a window in the ai session, launch claude, return window id.
create_window() {
  local dir="$1" name="$2" wid instructions
  wid=$(tmux new-window -a -t "$SESSION:{end}" -n "$name" -c "$dir" -P -F '#{window_id}')
  sleep 0.1
  instructions=$(build_instructions "$name")
  tmux send-keys -t "$wid" "${LAUNCH_PREFIX}$AGENT_CMD \"$instructions\"" Enter
  echo "$wid"
}

# Add a directory to candidates only if it's a Claude-template project.
maybe_add_candidate() {
  local dir="$1"
  [ -f "$dir/.ai/protocols.org" ] && candidates+=("~/${dir#"$HOME"/}")
}

# Build un-annotated candidate list (used by multi_mode and sort_windows).
build_candidates() {
  candidates=()
  maybe_add_candidate "$HOME/.emacs.d"
  maybe_add_candidate "$HOME/.dotfiles"
  if [ -d "$HOME/code" ]; then
    while IFS= read -r d; do
      maybe_add_candidate "$d"
    done < <(find "$HOME/code" -maxdepth 1 -mindepth 1 -type d | sort)
  fi
  if [ -d "$HOME/projects" ]; then
    while IFS= read -r d; do
      maybe_add_candidate "$d"
    done < <(find "$HOME/projects" -maxdepth 1 -mindepth 1 -type d | sort)
  fi
}

# Parallel fetch across all current candidates (capped concurrency).
fetch_candidates() {
  local max=6 running=0 dir
  for c in "${candidates[@]}"; do
    dir="${c/#\~/$HOME}"
    [ -d "$dir/.git" ] || continue
    git -C "$dir" fetch --quiet 2>/dev/null &
    running=$((running + 1))
    if [ "$running" -ge "$max" ]; then
      wait
      running=0
    fi
  done
  wait
}

# Return " (↑N ↓N dirty)" or " (✓)" if clean.
git_status_indicator() {
  local dir="$1" upstream ahead=0 behind=0 parts=()
  [ -d "$dir/.git" ] || return 0

  upstream=$(git -C "$dir" rev-parse --abbrev-ref --symbolic-full-name @{u} 2>/dev/null || true)
  if [ -n "$upstream" ]; then
    ahead=$(git -C "$dir" rev-list --count "$upstream..HEAD" 2>/dev/null || echo 0)
    behind=$(git -C "$dir" rev-list --count "HEAD..$upstream" 2>/dev/null || echo 0)
    [ "${ahead:-0}" -gt 0 ] 2>/dev/null && parts+=("↑$ahead")
    [ "${behind:-0}" -gt 0 ] 2>/dev/null && parts+=("↓$behind")
  else
    parts+=("no upstream")
  fi

  if ! git -C "$dir" diff --quiet 2>/dev/null \
    || ! git -C "$dir" diff --cached --quiet 2>/dev/null \
    || [ -n "$(git -C "$dir" ls-files --others --exclude-standard 2>/dev/null)" ]; then
    parts+=("dirty")
  fi

  if [ ${#parts[@]} -gt 0 ]; then
    local IFS=' '
    printf ' (%s)' "${parts[*]}"
  else
    printf ' (✓)'
  fi
}

# Append status annotations to every candidate.
annotate_candidates() {
  local dir status annotated=()
  for c in "${candidates[@]}"; do
    dir="${c/#\~/$HOME}"
    status=$(git_status_indicator "$dir")
    annotated+=("${c}${status}")
  done
  candidates=("${annotated[@]}")
}

# Pull if clean, behind, not ahead. No-op otherwise.
auto_pull_if_clean() {
  local dir="$1" upstream ahead behind
  [ -d "$dir/.git" ] || return 0

  if ! git -C "$dir" diff --quiet 2>/dev/null \
    || ! git -C "$dir" diff --cached --quiet 2>/dev/null \
    || [ -n "$(git -C "$dir" ls-files --others --exclude-standard 2>/dev/null)" ]; then
    return 0
  fi

  upstream=$(git -C "$dir" rev-parse --abbrev-ref --symbolic-full-name @{u} 2>/dev/null || true)
  [ -z "$upstream" ] && return 0

  ahead=$(git -C "$dir" rev-list --count "$upstream..HEAD" 2>/dev/null || echo 0)
  [ "${ahead:-0}" -gt 0 ] 2>/dev/null && return 0

  behind=$(git -C "$dir" rev-list --count "HEAD..$upstream" 2>/dev/null || echo 0)
  [ "${behind:-0}" -eq 0 ] 2>/dev/null && return 0

  git -C "$dir" pull --ff-only --quiet 2>/dev/null || true
}

# Strip " (annotation)" suffix from fzf output so downstream gets raw paths.
read_selections() {
  selected=()
  while IFS= read -r line; do
    selected+=("${line%% (*}")
  done <<<"$1"
}

# Re-order windows: non-project windows at base-index, projects alphabetically after.
sort_windows() {
  local windows others="" projects="" base_idx project_names=""
  base_idx=$(tmux show-option -gv base-index 2>/dev/null || echo 0)
  windows=$(tmux list-windows -t "$SESSION" -F '#{window_name}'$'\t''#{window_id}')

  build_candidates
  for c in "${candidates[@]}"; do
    project_names+="$(basename "${c/#\~/$HOME}")"$'\n'
  done

  while IFS=$'\t' read -r wname wid; do
    [ -z "$wname" ] && continue
    if echo "$project_names" | grep -qxF "$wname"; then
      projects+="${wname}"$'\t'"${wid}"$'\n'
    else
      others+="${wname}"$'\t'"${wid}"$'\n'
    fi
  done <<<"$windows"
  others=$(echo -n "$others" | sort -t$'\t' -k1,1f)
  projects=$(echo -n "$projects" | sort -t$'\t' -k1,1f)

  local all
  all=$(printf '%s\n' "$others" "$projects" | sed '/^$/d')

  local i=900
  while IFS=$'\t' read -r _n wid; do
    tmux move-window -s "$wid" -t "$SESSION:$i"
    i=$((i + 1))
  done <<<"$all"

  i=$base_idx
  if [ -n "$others" ]; then
    while IFS=$'\t' read -r _n wid; do
      tmux move-window -s "$wid" -t "$SESSION:$i"
      i=$((i + 1))
    done <<<"$others"
  fi
  if [ -n "$projects" ]; then
    while IFS=$'\t' read -r _n wid; do
      tmux move-window -s "$wid" -t "$SESSION:$i"
      i=$((i + 1))
    done <<<"$projects"
  fi
}

# Find existing window id in ai session by window name; empty if none.
find_window_id() {
  local name="$1"
  tmux list-windows -t "$SESSION" -F '#{window_name}'$'\t''#{window_id}' 2>/dev/null \
    | awk -F'\t' -v n="$name" '$1 == n {print $2; exit}'
}

# Git prep for a single directory. Uses FETCH_HEAD cache to skip back-to-back
# fetches. Pulls automatically if clean-and-behind; prints one-line summary
# if diverged/dirty/ahead.
prep_git_single() {
  local dir="$1" gitdir upstream ahead=0 behind=0 dirty="" age fetch_stale=1 parts=()
  git -C "$dir" rev-parse --is-inside-work-tree >/dev/null 2>&1 || return 0

  gitdir=$(git -C "$dir" rev-parse --git-dir 2>/dev/null)
  if [ -f "$gitdir/FETCH_HEAD" ]; then
    age=$(( $(date +%s) - $(stat -c %Y "$gitdir/FETCH_HEAD" 2>/dev/null || echo 0) ))
    [ "$age" -lt 600 ] && fetch_stale=0
  fi
  [ "$fetch_stale" -eq 1 ] && git -C "$dir" fetch --quiet 2>/dev/null || true

  upstream=$(git -C "$dir" rev-parse --abbrev-ref --symbolic-full-name @{u} 2>/dev/null || true)
  [ -z "$upstream" ] && return 0

  ahead=$(git -C "$dir" rev-list --count "$upstream..HEAD" 2>/dev/null || echo 0)
  behind=$(git -C "$dir" rev-list --count "HEAD..$upstream" 2>/dev/null || echo 0)

  if ! git -C "$dir" diff --quiet 2>/dev/null \
    || ! git -C "$dir" diff --cached --quiet 2>/dev/null \
    || [ -n "$(git -C "$dir" ls-files --others --exclude-standard 2>/dev/null)" ]; then
    dirty="dirty"
  fi

  if [ -z "$dirty" ] && [ "${ahead:-0}" -eq 0 ] && [ "${behind:-0}" -gt 0 ]; then
    echo "ai: pulling $behind commit(s) from $upstream..." >&2
    git -C "$dir" pull --ff-only --quiet
  elif [ "${ahead:-0}" -gt 0 ] || [ "${behind:-0}" -gt 0 ] || [ -n "$dirty" ]; then
    [ "${ahead:-0}" -gt 0 ] && parts+=("↑$ahead")
    [ "${behind:-0}" -gt 0 ] && parts+=("↓$behind")
    [ -n "$dirty" ] && parts+=("$dirty")
    echo "ai: $(basename "$dir") — ${parts[*]}" >&2
  fi
}

# ---------- modes ----------

attach_mode() {
  if ! tmux has-session -t "$SESSION" 2>/dev/null; then
    echo "ai: no '$SESSION' session to attach to" >&2
    exit 1
  fi
  sort_windows
  attach_session
}

# Open a single project (or focus existing window).
single_mode() {
  local arg="$1" dir name wid existing
  dir="$(cd "$arg" 2>/dev/null && pwd)" || { echo "ai: cannot access '$arg'" >&2; return 1; }

  if [ ! -f "$dir/.ai/protocols.org" ]; then
    echo "ai: $dir has no .ai/protocols.org — not a Claude-template project" >&2
    return 1
  fi

  name="$(basename "$dir")"

  # Window already exists? Focus it, skip git prep (claude is running in there).
  if tmux has-session -t "$SESSION" 2>/dev/null; then
    existing=$(find_window_id "$name")
    if [ -n "$existing" ]; then
      tmux select-window -t "$existing"
      attach_session
      return 0
    fi
  fi

  prep_git_single "$dir"

  # Create session with first window, or add a new window.
  if tmux has-session -t "$SESSION" 2>/dev/null; then
    wid=$(create_window "$dir" "$name")
  else
    local instructions
    wid=$(tmux new-session -d -s "$SESSION" -n "$name" -c "$dir" -P -F '#{window_id}')
    instructions=$(build_instructions "$name")
    tmux send-keys -t "$wid" "${LAUNCH_PREFIX}$AGENT_CMD \"$instructions\"" Enter
  fi

  sort_windows
  tmux select-window -t "$wid"
  attach_session
}

# Multi-select via fzf (the original aix flow).
multi_mode() {
  local filtered=() selections first_wid=""

  if tmux has-session -t "$SESSION" 2>/dev/null; then
    # Add to existing session — filter out projects already open.
    local existing
    existing=$(tmux list-windows -t "$SESSION" -F '#{window_name}')
    build_candidates
    for c in "${candidates[@]}"; do
      local n
      n="$(basename "${c/#\~/$HOME}")"
      if ! echo "$existing" | grep -qxF "$n"; then
        filtered+=("$c")
      fi
    done
    if [ ${#filtered[@]} -eq 0 ]; then
      echo "All projects already have windows open."
      sort_windows
      attach_session
      return 0
    fi
    candidates=("${filtered[@]}")
  else
    build_candidates
  fi

  echo "Fetching remotes..." >&2
  fetch_candidates
  annotate_candidates

  # One-line summary: total / clean / changed
  local _total=${#candidates[@]} _clean _changed
  _clean=$(printf '%s\n' "${candidates[@]}" | grep -c ' (✓)$' || true)
  _changed=$((_total - _clean))
  echo "Fetched $_total projects: $_clean clean, $_changed with changes" >&2

  selections=$(printf '%s\n' "${candidates[@]}" | fzf --multi --height=70% --reverse) || true
  [ -z "$selections" ] && return 0

  read_selections "$selections"

  if ! tmux has-session -t "$SESSION" 2>/dev/null; then
    # Create session with first selection
    local first="${selected[0]}"
    local dir name
    dir="${first/#\~/$HOME}"
    name="$(basename "$dir")"
    auto_pull_if_clean "$dir"
    local instructions
    first_wid=$(tmux new-session -d -s "$SESSION" -n "$name" -c "$dir" -P -F '#{window_id}')
    instructions=$(build_instructions "$name")
    tmux send-keys -t "$first_wid" "${LAUNCH_PREFIX}$AGENT_CMD \"$instructions\"" Enter
    for entry in "${selected[@]:1}"; do
      dir="${entry/#\~/$HOME}"
      name="$(basename "$dir")"
      auto_pull_if_clean "$dir"
      create_window "$dir" "$name" > /dev/null
    done
  else
    # Add windows to existing session
    for entry in "${selected[@]}"; do
      local dir name wid
      dir="${entry/#\~/$HOME}"
      name="$(basename "$dir")"
      auto_pull_if_clean "$dir"
      wid=$(create_window "$dir" "$name")
      [ -z "$first_wid" ] && first_wid="$wid"
    done
  fi

  sort_windows
  [ -n "$first_wid" ] && tmux select-window -t "$first_wid"
  attach_session
}

# Print the launch command a real run would send to the pane, then exit.
# Exists for the launcher's bats tests: exercises runtime resolution and the
# opening line with no tmux or fzf involved.
print_launch_mode() {
  local arg="$1" dir name
  dir="$(cd "$arg" 2>/dev/null && pwd)" || { echo "ai: cannot access '$arg'" >&2; exit 1; }
  if [ ! -f "$dir/.ai/protocols.org" ]; then
    echo "ai: $dir has no .ai/protocols.org — not an agent-template project" >&2
    exit 1
  fi
  name="$(basename "$dir")"
  printf '%s "%s"\n' "$AGENT_CMD" "$(build_instructions "$name")"
  exit 0
}

# ---------- dispatch ----------

print_launch=""
while [ $# -gt 0 ]; do
  case "$1" in
    -h|--help)
      usage
      ;;
    --runtime)
      [ -z "${2:-}" ] && { echo "ai: --runtime needs a value — valid runtimes: claude, codex" >&2; exit 2; }
      RUNTIME="$2"
      shift 2
      ;;
    --runtime=*)
      RUNTIME="${1#--runtime=}"
      shift
      ;;
    --print-launch)
      print_launch=1
      shift
      ;;
    *)
      break
      ;;
  esac
done

resolve_agent_cmd

if [ -n "$print_launch" ]; then
  [ $# -eq 0 ] && { echo "ai: --print-launch needs a project directory" >&2; exit 2; }
  print_launch_mode "$1"
fi

check_deps

case "${1:-}" in
  --attach)
    attach_mode
    ;;
  "")
    multi_mode
    ;;
  *)
    for arg in "$@"; do
      single_mode "$arg"
    done
    ;;
esac