#!/bin/bash
# Launch tmux session with an LLM in fzf-selected project directories

SESSION="ai"

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

# Dependency checks
for cmd in fzf tmux claude; do
  if ! command -v "$cmd" &>/dev/null; then
    echo "Error: $cmd is not installed" >&2
    exit 1
  fi
done

# Claude command — separated for safe editing
AI_CMD="claude"
AI_INSTRUCTIONS='Read .ai/protocols.org and follow all instructions.'

# Create a window in the session and launch claude; prints window ID
create_window() {
  local dir="$1" name="$2" wid
  wid=$(tmux new-window -t "$SESSION" -n "$name" -c "$dir" -P -F '#{window_id}')
  sleep 0.1
  tmux send-keys -t "$wid" "$AI_CMD \"$AI_INSTRUCTIONS\"" Enter
  echo "$wid"
}

# Read fzf selections into the 'selected' array — strip any " (annotation)"
# suffix so downstream code sees the raw tilde path.
read_selections() {
  selected=()
  while IFS= read -r line; do
    selected+=("${line%% (*}")
  done <<<"$1"
}

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

# Build candidate directory list
build_candidates() {
  candidates=()
  maybe_add_candidate "$HOME/.emacs.d"
  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
}

# Fetch all candidate repos in parallel (capped concurrency).
# Blocking call — synchronous relative to the rest of the script.
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++))
    if ((running >= max)); then
      wait
      running=0
    fi
  done
  wait
}

# Produce a git-status indicator string for a directory, e.g.
# "(↑1 ↓3 dirty)" or "(no upstream)". Empty string if clean.
git_status_indicator() {
  local dir="$1" upstream ahead behind parts=()
  [ -d "$dir/.git" ] || return

  upstream=$(git -C "$dir" rev-parse --abbrev-ref --symbolic-full-name @{u} 2>/dev/null)
  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[*]}"
  fi
}

# Annotate the 'candidates' array with git-status indicators
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[@]}")
}

# Auto-pull a directory if its working tree is clean, upstream is set,
# the branch is behind (not ahead, not diverged). No-op otherwise.
auto_pull_if_clean() {
  local dir="$1" upstream ahead behind
  [ -d "$dir/.git" ] || return

  # Dirty? Skip.
  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
  fi

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

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

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

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

# Sort 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 list of known project basenames from candidate directories
  build_candidates
  project_names=""
  for c in "${candidates[@]}"; do
    project_names+="$(basename "${c/#\~/$HOME}")"$'\n'
  done

  # Classify: windows matching a project basename are projects, rest are non-project
  others=""
  projects=""
  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')

  # Move all to high temp indices to avoid collisions
  local i=900
  while IFS=$'\t' read -r _name wid; do
    tmux move-window -s "$wid" -t "$SESSION:$i"
    ((i++))
  done <<<"$all"

  # Place non-project windows at base-index, projects contiguously after
  i=$base_idx
  if [ -n "$others" ]; then
    while IFS=$'\t' read -r _name wid; do
      tmux move-window -s "$wid" -t "$SESSION:$i"
      ((i++))
    done <<<"$others"
  fi
  if [ -n "$projects" ]; then
    while IFS=$'\t' read -r _name wid; do
      tmux move-window -s "$wid" -t "$SESSION:$i"
      ((i++))
    done <<<"$projects"
  fi
}

# Attach directly to existing session
if [ "$1" = "--attach" ]; then
  if ! tmux has-session -t "$SESSION" 2>/dev/null; then
    echo "No $SESSION session to attach to." >&2
    exit 1
  fi
  sort_windows
  attach_session
  exit 0
fi

# If session exists, add new windows to it
if tmux has-session -t "$SESSION" 2>/dev/null; then
  # Get existing window names to filter duplicates
  existing=$(tmux list-windows -t "$SESSION" -F '#{window_name}')

  build_candidates

  # Filter out candidates that already have a window open
  filtered=()
  for c in "${candidates[@]}"; do
    name="$(basename "${c/#\~/$HOME}")"
    if ! echo "$existing" | grep -qxF "$name"; then
      filtered+=("$c")
    fi
  done

  if [ ${#filtered[@]} -eq 0 ]; then
    echo "All projects already have windows open."
  else
    # Annotate filtered list with git-status indicators (fetch first)
    candidates=("${filtered[@]}")
    echo "Fetching remotes..." >&2
    fetch_candidates
    annotate_candidates

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

    if [ -n "$selections" ]; then
      read_selections "$selections"

      first_wid=""
      for entry in "${selected[@]}"; do
        dir="${entry/#\~/$HOME}"
        name="$(basename "$dir")"
        auto_pull_if_clean "$dir"
        wid=$(create_window "$dir" "$name")
        [ -z "$first_wid" ] && first_wid="$wid"
      done
    fi
  fi

  sort_windows
  [ -n "$first_wid" ] && tmux select-window -t "$first_wid"
  attach_session
  exit 0
fi

# New session: select projects and create session
build_candidates

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

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

read_selections "$selections"

# Create session with first selection
first="${selected[0]}"
dir="${first/#\~/$HOME}"
name="$(basename "$dir")"
auto_pull_if_clean "$dir"
first_wid=$(tmux new-session -d -s "$SESSION" -n "$name" -c "$dir" -P -F '#{window_id}')
tmux send-keys -t "$first_wid" "$AI_CMD \"$AI_INSTRUCTIONS\"" Enter

# Create remaining windows
for entry in "${selected[@]:1}"; do
  dir="${entry/#\~/$HOME}"
  name="$(basename "$dir")"
  auto_pull_if_clean "$dir"
  create_window "$dir" "$name" > /dev/null
done

# Sort windows and select first project
sort_windows
tmux select-window -t "$first_wid"
attach_session