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

usage() {
  cat <<'EOF'
Create one jj workspace per issue, optionally create bookmarks, and optionally open zellij tabs running pi.

Usage:
  jj-workspace-fanout.sh [options] --issue "KEY=Title" [--issue "KEY=Title" ...]

Options:
  --repo PATH             Repo root (default: current directory)
  --base REV              Base revision/bookmark (default: master)
  --remote NAME           Git remote to fetch from (default: origin)
  --issue KEY=TITLE       Issue key and title (repeatable)
  --session NAME          Zellij session name (defaults to ZELLIJ_SESSION_NAME if set)
  --open-pi               Open a zellij tab per workspace and launch pi
  --no-fetch              Skip jj git fetch
  --no-bookmarks          Do not create feature/<issue> bookmarks
  --keep-existing         Skip creation for existing workspaces instead of failing
  --reset-existing        Forget and delete existing workspaces before recreating them
  --prompt-suffix TEXT    Extra text appended to each pi prompt
  --pi-cmd CMD            pi command to launch (default: pi)
  --dry-run               Print planned actions without making changes
  --help                  Show this help

Examples:
  jj-workspace-fanout.sh \
    --repo /path/to/Phoenix \
    --base master \
    --issue "SPA-748=Wrap text in credits line items" \
    --issue "SPA-754=Resize seat count picker"

  jj-workspace-fanout.sh \
    --repo /path/to/Phoenix \
    --base master \
    --session attio \
    --open-pi \
    --issue "SPA-748=Wrap text in credits line items"
EOF
}

require_cmd() {
  if ! command -v "$1" >/dev/null 2>&1; then
    echo "error: missing required command: $1" >&2
    exit 1
  fi
}

shell_escape() {
  printf '%q' "$1"
}

log() {
  printf '[jj-issue-workspaces] %s\n' "$*"
}

run() {
  if [[ "$DRY_RUN" -eq 1 ]]; then
    printf '[dry-run] '
    printf '%q ' "$@"
    printf '\n'
  else
    "$@"
  fi
}

workspace_exists() {
  local workspace_name="$1"
  jj -R "$REPO" workspace list | awk -F: '{print $1}' | grep -Fxq "$workspace_name"
}

bookmark_exists() {
  local workspace_dir="$1"
  local bookmark_name="$2"
  jj -R "$workspace_dir" bookmark list "$bookmark_name" 2>/dev/null | grep -Eq "^${bookmark_name}:"
}

close_tab_if_exists() {
  local session_name="$1"
  local tab_name="$2"
  local tabs

  tabs=$(zellij --session "$session_name" action query-tab-names 2>/dev/null || true)
  if printf '%s\n' "$tabs" | grep -Fxq "$tab_name"; then
    log "closing existing zellij tab $tab_name"
    run zellij --session "$session_name" action go-to-tab-name "$tab_name"
    run zellij --session "$session_name" action close-tab
  fi
}

launch_pi_tab() {
  local session_name="$1"
  local tab_name="$2"
  local workspace_dir="$3"
  local prompt="$4"
  local cmd

  cmd="cd $(shell_escape "$workspace_dir") && pwd && $PI_CMD $(shell_escape "$prompt")"

  close_tab_if_exists "$session_name" "$tab_name"
  run zellij --session "$session_name" action new-tab --name "$tab_name"
  run zellij --session "$session_name" action write-chars "$cmd"
  run zellij --session "$session_name" action write 10
}

REPO="$(pwd)"
BASE="master"
REMOTE="origin"
SESSION="${ZELLIJ_SESSION_NAME:-}"
OPEN_PI=0
FETCH=1
CREATE_BOOKMARKS=1
KEEP_EXISTING=0
RESET_EXISTING=0
DRY_RUN=0
PROMPT_SUFFIX=""
PI_CMD="pi"
declare -a ISSUES=()

while [[ $# -gt 0 ]]; do
  case "$1" in
    --repo)
      REPO="$2"
      shift 2
      ;;
    --base)
      BASE="$2"
      shift 2
      ;;
    --remote)
      REMOTE="$2"
      shift 2
      ;;
    --issue)
      ISSUES+=("$2")
      shift 2
      ;;
    --session)
      SESSION="$2"
      shift 2
      ;;
    --open-pi)
      OPEN_PI=1
      shift
      ;;
    --no-fetch)
      FETCH=0
      shift
      ;;
    --no-bookmarks)
      CREATE_BOOKMARKS=0
      shift
      ;;
    --keep-existing)
      KEEP_EXISTING=1
      shift
      ;;
    --reset-existing)
      RESET_EXISTING=1
      shift
      ;;
    --prompt-suffix)
      PROMPT_SUFFIX="$2"
      shift 2
      ;;
    --pi-cmd)
      PI_CMD="$2"
      shift 2
      ;;
    --dry-run)
      DRY_RUN=1
      shift
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      echo "error: unknown argument: $1" >&2
      usage >&2
      exit 1
      ;;
  esac
done

if [[ ${#ISSUES[@]} -eq 0 ]]; then
  echo "error: at least one --issue KEY=TITLE is required" >&2
  exit 1
fi

if [[ "$KEEP_EXISTING" -eq 1 && "$RESET_EXISTING" -eq 1 ]]; then
  echo "error: --keep-existing and --reset-existing cannot be combined" >&2
  exit 1
fi

REPO="$(cd "$REPO" && pwd)"
PARENT_DIR="$(dirname "$REPO")"
REPO_BASENAME="$(basename "$REPO")"

require_cmd jj
if [[ "$OPEN_PI" -eq 1 ]]; then
  require_cmd zellij
  if [[ -z "$SESSION" ]]; then
    echo "error: --open-pi requires --session <name> or ZELLIJ_SESSION_NAME" >&2
    exit 1
  fi
fi

if [[ ! -d "$REPO/.jj" ]]; then
  echo "error: repo is not a jj repository: $REPO" >&2
  exit 1
fi

if [[ "$FETCH" -eq 1 ]]; then
  log "fetching $BASE from $REMOTE"
  run jj -R "$REPO" git fetch --remote "$REMOTE" --branch "$BASE"
fi

log "validating base revision $BASE"
run jj -R "$REPO" log -r "$BASE" --no-pager

created_workspaces=()

for issue in "${ISSUES[@]}"; do
  if [[ "$issue" != *=* ]]; then
    echo "error: issue must be formatted as KEY=TITLE: $issue" >&2
    exit 1
  fi

  issue_key="${issue%%=*}"
  issue_title="${issue#*=}"
  issue_slug="$(printf '%s' "$issue_key" | tr '[:upper:]' '[:lower:]')"
  workspace_name="$issue_slug"
  workspace_dir="$PARENT_DIR/${REPO_BASENAME}-${issue_slug}"
  bookmark_name="feature/$issue_slug"
  prompt="Work on ${issue_key}: ${issue_title}. You are in the dedicated jj workspace for this issue. First inspect the relevant code, identify the main components involved, and propose a short plan before editing."

  if [[ -n "$PROMPT_SUFFIX" ]]; then
    prompt+=" ${PROMPT_SUFFIX}"
  fi

  if workspace_exists "$workspace_name" || [[ -e "$workspace_dir" ]]; then
    if [[ "$RESET_EXISTING" -eq 1 ]]; then
      log "resetting existing workspace $workspace_name"
      if workspace_exists "$workspace_name"; then
        run jj -R "$REPO" workspace forget "$workspace_name"
      fi
      run rm -rf "$workspace_dir"
    elif [[ "$KEEP_EXISTING" -eq 1 ]]; then
      log "keeping existing workspace $workspace_name at $workspace_dir"
    else
      echo "error: workspace already exists: $workspace_name ($workspace_dir). Use --keep-existing or --reset-existing." >&2
      exit 1
    fi
  fi

  if ! workspace_exists "$workspace_name"; then
    log "creating workspace $workspace_name at $workspace_dir"
    run jj -R "$REPO" workspace add --name "$workspace_name" -r "$BASE" "$workspace_dir"
  fi

  if [[ "$CREATE_BOOKMARKS" -eq 1 ]]; then
    log "ensuring bookmark $bookmark_name exists"
    if bookmark_exists "$workspace_dir" "$bookmark_name"; then
      run jj -R "$workspace_dir" bookmark set "$bookmark_name" -r @
    else
      run jj -R "$workspace_dir" bookmark create "$bookmark_name"
    fi
  fi

  if [[ "$OPEN_PI" -eq 1 ]]; then
    log "opening zellij tab $workspace_name in session $SESSION"
    run launch_pi_tab "$SESSION" "$workspace_name" "$workspace_dir" "$prompt"
  fi

  created_workspaces+=("$workspace_name:$workspace_dir:$bookmark_name")
done

printf '\nCreated/updated workspaces:\n'
for item in "${created_workspaces[@]}"; do
  IFS=':' read -r workspace_name workspace_dir bookmark_name <<<"$item"
  printf '  - %s -> %s' "$workspace_name" "$workspace_dir"
  if [[ "$CREATE_BOOKMARKS" -eq 1 ]]; then
    printf ' [%s]' "$bookmark_name"
  fi
  printf '\n'
done

if [[ "$OPEN_PI" -eq 1 ]]; then
  printf '\nZellij session: %s\n' "$SESSION"
fi