#!/usr/bin/env bash

set -euo pipefail

# This script safely imports selected upstream commits onto the current local
# branch through cherry-pick.
# It:
# 1. checks that no conflicting git operation is already in progress;
# 2. ensures the upstream remote exists and points to the official repository;
# 3. lists upstream commits whose changes are not yet present on the current
#    branch;
# 4. creates a safety backup branch and an exact worktree snapshot;
# 5. stashes tracked and untracked local changes before the import;
# 6. cherry-picks the selected upstream commits with `-x` so the original
#    upstream commit hash remains visible in the imported commit message;
# 7. restores the stashed local changes only after the import completes.
# If a conflict happens, the backup branch, the exact snapshot, and the saved
# stash are all kept so the original state can be restored safely.

BRANCH="${BRANCH:-main}"
REMOTE_NAME="${REMOTE_NAME:-upstream}"
REMOTE_URL="${REMOTE_URL:-https://github.com/go-gitea/gitea.git}"
REPO_ROOT=""
GIT_DIR=""
STATE_FILE=""
SNAPSHOT_ROOT=""
CURRENT_BRANCH=""
BACKUP_BRANCH=""
STASH_REF=""
STASH_HASH=""
STATE_BRANCH=""
STATE_BACKUP_BRANCH=""
STATE_STASH_HASH=""
STATE_REMOTE_NAME=""
STATE_REMOTE_URL=""
STATE_TARGET_BRANCH=""
STATE_START_HEAD=""
STATE_SNAPSHOT_PATH=""
STATE_IMPORT_COMMITS=""
STATE_STATUS=""
STATE_UPDATED_AT=""

say() {
  printf '%s\n' "$*"
}

die() {
  say "ERROR: $*" >&2
  exit 1
}

show_usage() {
  cat <<EOF
Usage: $(basename "$0") [command]

Available commands:
  fetch
          Fetch the latest upstream changes only.
  list
          Show upstream commits whose changes are not yet present on the current branch.
  import <commit> [commit...]
          Cherry-pick one or more selected upstream commits with \`-x\`.
  import-range <from-commit> <to-commit>
          Cherry-pick an inclusive upstream commit range in upstream order.
  continue
          Continue an interrupted cherry-pick and restore the saved local stash once the import finishes.
  backup
          Create a safety backup branch from the current HEAD only.
  rollback [backup-branch|snapshot-dir]
          Restore the saved pre-import state, or an explicit backup branch / snapshot.
  restore-stash
          Re-apply the saved pre-import local stash from the last recorded import state.
  restore-snapshot [snapshot-dir]
          Restore the exact saved repository snapshot from the last recorded real action, or from a specific snapshot directory.
  status
          Show the current repository, branch, remote, local-change state, and saved import metadata.
  help
          Show this help text.

If no command is provided and the script is run from a terminal, an interactive
menu is shown.

Recommended order:
  1. Run:  ./.import-upstream-cherry-pick.sh status
  2. Run:  ./.import-upstream-cherry-pick.sh list
  3. Choose the exact upstream commits you want to import.
  4. Run:  ./.import-upstream-cherry-pick.sh import <commit...>
     or:   ./.import-upstream-cherry-pick.sh import-range <from> <to>
  5. If cherry-pick stops with conflicts:
     - either resolve them and continue with: ./.import-upstream-cherry-pick.sh continue
     - or return safely with: ./.import-upstream-cherry-pick.sh rollback
  6. If restoring local changes after the import causes conflicts:
     - either resolve them manually
     - or return safely with: ./.import-upstream-cherry-pick.sh rollback
  7. Use:  ./.import-upstream-cherry-pick.sh restore-snapshot
     when you want to restore the exact saved repository state, including
     untracked and ignored files that existed at snapshot time.
EOF
}

ensure_git_repo() {
  local raw_git_dir

  REPO_ROOT="$(git rev-parse --show-toplevel 2>/dev/null)" || die "This script must be run inside a git repository."
  raw_git_dir="$(git rev-parse --git-dir)"
  case "$raw_git_dir" in
    /*) GIT_DIR="$raw_git_dir" ;;
    *) GIT_DIR="$REPO_ROOT/$raw_git_dir" ;;
  esac
  STATE_FILE="$GIT_DIR/.import-upstream-cherry-pick-state"
  SNAPSHOT_ROOT="$GIT_DIR/.import-upstream-cherry-pick-snapshots"
}

write_state() {
  {
    printf 'STATE_BRANCH=%q\n' "$STATE_BRANCH"
    printf 'STATE_BACKUP_BRANCH=%q\n' "$STATE_BACKUP_BRANCH"
    printf 'STATE_STASH_HASH=%q\n' "$STATE_STASH_HASH"
    printf 'STATE_REMOTE_NAME=%q\n' "$STATE_REMOTE_NAME"
    printf 'STATE_REMOTE_URL=%q\n' "$STATE_REMOTE_URL"
    printf 'STATE_TARGET_BRANCH=%q\n' "$STATE_TARGET_BRANCH"
    printf 'STATE_START_HEAD=%q\n' "$STATE_START_HEAD"
    printf 'STATE_SNAPSHOT_PATH=%q\n' "$STATE_SNAPSHOT_PATH"
    printf 'STATE_IMPORT_COMMITS=%q\n' "$STATE_IMPORT_COMMITS"
    printf 'STATE_STATUS=%q\n' "$STATE_STATUS"
    printf 'STATE_UPDATED_AT=%q\n' "$(date +%Y-%m-%dT%H:%M:%S)"
  } >"$STATE_FILE"
}

load_state() {
  ensure_git_repo
  [ -f "$STATE_FILE" ] || return 1
  # shellcheck disable=SC1090
  . "$STATE_FILE"
  return 0
}

worktree_has_local_changes() {
  ! git diff --quiet || ! git diff --cached --quiet || [ -n "$(git ls-files --others --exclude-standard)" ]
}

has_cherry_pick_in_progress() {
  [ -e "$(git rev-parse --git-path CHERRY_PICK_HEAD)" ]
}

sanitize_ref_name() {
  printf '%s\n' "${1//\//-}"
}

ensure_rsync_available() {
  command -v rsync >/dev/null 2>&1 || die "rsync is required for exact repository snapshots."
}

snapshot_repo_dir() {
  printf '%s\n' "$1/repo"
}

is_valid_snapshot_dir() {
  [ -d "$1" ] && [ -d "$(snapshot_repo_dir "$1")" ]
}

snapshot_metadata_value() {
  local snapshot_dir="$1" key="$2"

  [ -f "$snapshot_dir/metadata.txt" ] || return 1
  sed -n "s/^${key}=//p" "$snapshot_dir/metadata.txt" | head -n 1
}

create_worktree_snapshot_for_reason() {
  local reason="$1" safe_branch_name timestamp snapshot_dir snapshot_dir_repo

  ensure_rsync_available
  safe_branch_name="$(sanitize_ref_name "$CURRENT_BRANCH")"
  timestamp="$(date +%Y%m%d-%H%M%S)"
  snapshot_dir="$SNAPSHOT_ROOT/$safe_branch_name/${timestamp}-${reason}"
  snapshot_dir_repo="$(snapshot_repo_dir "$snapshot_dir")"

  mkdir -p "$snapshot_dir_repo"
  rsync --archive --delete --exclude='.git' "$REPO_ROOT"/ "$snapshot_dir_repo"/

  {
    printf 'branch=%s\n' "$CURRENT_BRANCH"
    printf 'head=%s\n' "$(git rev-parse HEAD)"
    printf 'created_at=%s\n' "$(date +%Y-%m-%dT%H:%M:%S)"
    printf 'reason=%s\n' "$reason"
  } >"$snapshot_dir/metadata.txt"

  printf '%s\n' "$snapshot_dir"
}

restore_worktree_snapshot() {
  local snapshot_dir="$1" snapshot_dir_repo

  ensure_rsync_available
  is_valid_snapshot_dir "$snapshot_dir" || die "Snapshot directory '$snapshot_dir' is not valid."
  snapshot_dir_repo="$(snapshot_repo_dir "$snapshot_dir")"

  rsync --archive --delete --exclude='.git' "$snapshot_dir_repo"/ "$REPO_ROOT"/
}

find_latest_snapshot_for_current() {
  local safe_branch_name snapshot_branch_dir
  safe_branch_name="$(sanitize_ref_name "$CURRENT_BRANCH")"
  snapshot_branch_dir="$SNAPSHOT_ROOT/$safe_branch_name"
  [ -d "$snapshot_branch_dir" ] || return 1
  find "$snapshot_branch_dir" -mindepth 1 -maxdepth 1 -type d | sort -r | head -n 1
}

ensure_no_noncherrypick_git_operation_in_progress() {
  local git_path
  for git_path in MERGE_HEAD REVERT_HEAD BISECT_LOG rebase-merge rebase-apply; do
    if [ -e "$(git rev-parse --git-path "$git_path")" ]; then
      die "A git operation is already in progress ($git_path). Finish it before running this script."
    fi
  done
}

ensure_no_git_operation_in_progress() {
  local git_path
  for git_path in CHERRY_PICK_HEAD MERGE_HEAD REVERT_HEAD BISECT_LOG rebase-merge rebase-apply; do
    if [ -e "$(git rev-parse --git-path "$git_path")" ]; then
      die "A git operation is already in progress ($git_path). Finish it before running this script."
    fi
  done

  if [ -n "$(git diff --name-only --diff-filter=U)" ]; then
    die "There are unmerged files in the working tree. Resolve them first."
  fi
}

ensure_current_branch() {
  CURRENT_BRANCH="$(git symbolic-ref --quiet --short HEAD || true)"
  [ -n "$CURRENT_BRANCH" ] || die "Detached HEAD is not supported. Check out a branch first."
}

ensure_current_branch_with_state_fallback() {
  CURRENT_BRANCH="$(git symbolic-ref --quiet --short HEAD || true)"
  if [ -n "$CURRENT_BRANCH" ]; then
    return
  fi

  if load_state && [ -n "${STATE_BRANCH:-}" ]; then
    CURRENT_BRANCH="$STATE_BRANCH"
    return
  fi

  die "Detached HEAD is not supported here and no saved branch state was found."
}

ensure_upstream_remote() {
  local current_remote_url

  if current_remote_url="$(git remote get-url "$REMOTE_NAME" 2>/dev/null)"; then
    if [ "$current_remote_url" != "$REMOTE_URL" ]; then
      die "Remote '$REMOTE_NAME' points to '$current_remote_url', expected '$REMOTE_URL'."
    fi
    return
  fi

  say "Adding remote '$REMOTE_NAME'..."
  git remote add "$REMOTE_NAME" "$REMOTE_URL"
}

create_backup_branch_for_reason() {
  local reason="$1" safe_branch_name timestamp branch_name

  safe_branch_name="$(sanitize_ref_name "$CURRENT_BRANCH")"
  timestamp="$(date +%Y%m%d-%H%M%S)"
  branch_name="backup/cherry-pick-${safe_branch_name}-${reason}-${timestamp}"

  git branch "$branch_name" HEAD >/dev/null
  printf '%s\n' "$branch_name"
}

find_latest_backup_branch_for_current() {
  local safe_branch_name

  safe_branch_name="$(sanitize_ref_name "$CURRENT_BRANCH")"
  git for-each-ref --sort=-creatordate --format='%(refname:short)' "refs/heads/backup/cherry-pick-${safe_branch_name}-*" | head -n 1
}

prepare_safety_before_real_action() {
  local backup_reason="$1"

  STATE_START_HEAD="$(git rev-parse HEAD)"
  BACKUP_BRANCH="$(create_backup_branch_for_reason "$backup_reason")"
  STATE_SNAPSHOT_PATH="$(create_worktree_snapshot_for_reason "$backup_reason")"
  say "Safety backup branch created: $BACKUP_BRANCH"
  say "Exact worktree snapshot created: $STATE_SNAPSHOT_PATH"
}

stash_current_changes() {
  local stash_prefix="$1" timestamp stash_name

  STASH_REF=""
  STASH_HASH=""
  if ! worktree_has_local_changes; then
    say "No local tracked/untracked changes to stash."
    return
  fi

  timestamp="$(date +%Y%m%d-%H%M%S)"
  stash_name="${stash_prefix}-$(sanitize_ref_name "$CURRENT_BRANCH")-${timestamp}"

  say "Stashing local tracked and untracked changes..."
  git stash push --include-untracked --message "$stash_name" >/dev/null
  STASH_REF="$(git stash list -1 --format='%gd')"
  [ -n "$STASH_REF" ] || die "Failed to locate the created stash entry."
  STASH_HASH="$(git rev-parse "$STASH_REF")"
  [ -n "$STASH_HASH" ] || die "Failed to resolve the created stash hash."
  say "Local changes saved in $STASH_REF"
}

fetch_upstream() {
  say "Fetching latest changes from $REMOTE_NAME/$BRANCH..."
  git fetch --prune "$REMOTE_NAME" "+refs/heads/$BRANCH:refs/remotes/$REMOTE_NAME/$BRANCH"
  git show-ref --verify --quiet "refs/remotes/$REMOTE_NAME/$BRANCH" || die "Remote branch '$REMOTE_NAME/$BRANCH' was not found."
}

update_state() {
  STATE_BRANCH="$CURRENT_BRANCH"
  STATE_BACKUP_BRANCH="$BACKUP_BRANCH"
  STATE_STASH_HASH="$STASH_HASH"
  STATE_REMOTE_NAME="$REMOTE_NAME"
  STATE_REMOTE_URL="$REMOTE_URL"
  STATE_TARGET_BRANCH="$BRANCH"
  write_state
}

find_stash_ref_by_hash() {
  local target_hash="$1" hash ref

  while read -r hash ref; do
    if [ "$hash" = "$target_hash" ]; then
      printf '%s\n' "$ref"
      return 0
    fi
  done < <(git stash list --format='%H %gd')
  return 1
}

drop_stash_by_hash_if_present() {
  local target_hash="$1" stash_ref

  [ -n "$target_hash" ] || return 0
  if stash_ref="$(find_stash_ref_by_hash "$target_hash")"; then
    git stash drop "$stash_ref" >/dev/null
  fi
}

restore_stash() {
  local keep_saved_copy="${1:-0}" restore_target

  if [ -z "$STASH_HASH" ] && [ -z "$STASH_REF" ]; then
    return
  fi

  restore_target="${STASH_HASH:-$STASH_REF}"
  say "Restoring stashed local changes from $restore_target..."
  if git stash apply --index "$restore_target"; then
    STASH_REF=""
    if [ "$keep_saved_copy" -eq 0 ]; then
      drop_stash_by_hash_if_present "$STASH_HASH"
      STASH_HASH=""
      STATE_STASH_HASH=""
    else
      STATE_STASH_HASH="${STASH_HASH:-$restore_target}"
    fi
    STATE_STATUS="completed"
    update_state
    if [ "$keep_saved_copy" -eq 0 ]; then
      say "Local changes restored successfully."
    else
      say "Local changes restored successfully. The original stash was kept for rollback safety."
    fi
    return
  fi

  STATE_STATUS="restore_conflict"
  update_state
  say "Conflicts occurred while restoring local changes."
  say "Backup branch kept at: $BACKUP_BRANCH"
  say "Stash kept at: $restore_target"
  say "Resolve the conflicts manually, or return safely with: ./.import-upstream-cherry-pick.sh rollback"
  exit 1
}

ordered_available_upstream_commits() {
  local available_commit_shas="" commit

  available_commit_shas="$(git cherry -v HEAD "$REMOTE_NAME/$BRANCH" | awk '$1=="+"{print $2}')"
  [ -n "$available_commit_shas" ] || return 0

  while IFS= read -r commit; do
    [ -n "$commit" ] || continue
    if printf '%s\n' "$available_commit_shas" | grep -Fxq "$commit"; then
      printf '%s\n' "$commit"
    fi
  done < <(git rev-list --reverse HEAD.."$REMOTE_NAME/$BRANCH")
}

commit_is_available_for_import() {
  local commit="$1"

  ordered_available_upstream_commits | grep -Fxq "$commit"
}

resolve_import_commit() {
  local commit="$1"

  git rev-parse --verify "${commit}^{commit}" 2>/dev/null || die "Commit '$commit' could not be resolved."
}

validate_importable_commit() {
  local commit="$1"

  git merge-base --is-ancestor "$commit" "$REMOTE_NAME/$BRANCH" || die "Commit '$commit' is not reachable from $REMOTE_NAME/$BRANCH."
  commit_is_available_for_import "$commit" || die "Commit '$commit' is not currently available for import into '$CURRENT_BRANCH'."
}

show_available_commits() {
  local found_any=0 commit

  while IFS= read -r commit; do
    [ -n "$commit" ] || continue
    found_any=1
    git show --no-patch --date=short --format='  - %h | %ad | %s' "$commit"
  done < <(ordered_available_upstream_commits)

  if [ "$found_any" -eq 0 ]; then
    say "No upstream commits are currently available for cherry-pick import."
  fi
}

run_status() {
  local stash_count latest_backup latest_snapshot

  ensure_git_repo
  ensure_current_branch_with_state_fallback
  ensure_upstream_remote

  stash_count="$(git stash list | wc -l | tr -d ' ')"
  latest_backup="$(find_latest_backup_branch_for_current || true)"
  latest_snapshot="$(find_latest_snapshot_for_current || true)"

  say "Repository: $(git rev-parse --show-toplevel)"
  say "Current branch: $CURRENT_BRANCH"
  say "Upstream remote: $REMOTE_NAME -> $REMOTE_URL"
  say "Working tree:"
  if ! worktree_has_local_changes; then
    say "  clean"
  else
    git status --short
  fi
  say "Stash entries: $stash_count"
  say "Latest backup for current branch: ${latest_backup:-none}"
  say "Latest snapshot for current branch: ${latest_snapshot:-none}"
  if load_state; then
    say "Saved import state: $STATE_STATUS"
    say "Saved backup branch: ${STATE_BACKUP_BRANCH:-none}"
    say "Saved starting commit: ${STATE_START_HEAD:-none}"
    say "Saved exact snapshot: ${STATE_SNAPSHOT_PATH:-none}"
    say "Saved original stash: ${STATE_STASH_HASH:-none}"
    say "Saved import commits: ${STATE_IMPORT_COMMITS:-none}"
  else
    say "Saved import state: none"
  fi
}

run_fetch_only() {
  ensure_git_repo
  ensure_no_git_operation_in_progress
  ensure_current_branch
  ensure_upstream_remote
  fetch_upstream
  say "Latest changes fetched from $REMOTE_NAME/$BRANCH."
}

run_list() {
  ensure_git_repo
  ensure_no_git_operation_in_progress
  ensure_current_branch
  ensure_upstream_remote
  fetch_upstream
  show_available_commits
}

run_backup_only() {
  ensure_git_repo
  ensure_no_git_operation_in_progress
  ensure_current_branch
  ensure_upstream_remote
  BACKUP_BRANCH="$(create_backup_branch_for_reason "manual-backup")"
  say "Safety backup branch created: $BACKUP_BRANCH"
}

run_import_commits() {
  local requested_commits=("$@") commit resolved_commit
  local -a selected_commits=()

  [ "${#requested_commits[@]}" -gt 0 ] || die "Provide at least one upstream commit to import."
  ensure_git_repo
  ensure_no_git_operation_in_progress
  ensure_current_branch
  ensure_upstream_remote
  fetch_upstream

  for commit in "${requested_commits[@]}"; do
    resolved_commit="$(resolve_import_commit "$commit")"
    validate_importable_commit "$resolved_commit"
    selected_commits+=("$resolved_commit")
  done

  STATE_IMPORT_COMMITS="${selected_commits[*]}"
  prepare_safety_before_real_action "before-cherry-pick-import"
  stash_current_changes "pre-cherry-pick-import"
  STATE_STATUS="prepared"
  update_state

  for commit in "${selected_commits[@]}"; do
    say "Cherry-picking upstream commit: $(git show --no-patch --format='%h %s' "$commit")"
    if ! git cherry-pick -x "$commit"; then
      STATE_STATUS="cherry_pick_conflict"
      update_state
      say "Cherry-pick stopped because of conflicts."
      say "Conflicting upstream commit: $(git show --no-patch --format='%h %s' CHERRY_PICK_HEAD)"
      say "Resolve conflicts and continue with: ./.import-upstream-cherry-pick.sh continue"
      say "Or return safely with: ./.import-upstream-cherry-pick.sh rollback"
      exit 1
    fi
  done

  restore_stash 1
  say "Imported ${#selected_commits[@]} upstream commit(s) by cherry-pick."
  say "Safety backup branch available at: $BACKUP_BRANCH"
  say "Safety snapshot available at: $STATE_SNAPSHOT_PATH"
  if [ -n "$STATE_STASH_HASH" ]; then
    say "Original pre-import stash retained for rollback at: $STATE_STASH_HASH"
  fi
}

run_import_range() {
  local from_commit="${1:-}" to_commit="${2:-}" from_resolved to_resolved commit
  local -a selected_commits=()

  [ -n "$from_commit" ] || die "Provide the start commit for import-range."
  [ -n "$to_commit" ] || die "Provide the end commit for import-range."
  ensure_git_repo
  ensure_no_git_operation_in_progress
  ensure_current_branch
  ensure_upstream_remote
  fetch_upstream

  from_resolved="$(resolve_import_commit "$from_commit")"
  to_resolved="$(resolve_import_commit "$to_commit")"
  git merge-base --is-ancestor "$from_resolved" "$REMOTE_NAME/$BRANCH" || die "Start commit '$from_commit' is not reachable from $REMOTE_NAME/$BRANCH."
  git merge-base --is-ancestor "$to_resolved" "$REMOTE_NAME/$BRANCH" || die "End commit '$to_commit' is not reachable from $REMOTE_NAME/$BRANCH."
  git merge-base --is-ancestor "$from_resolved" "$to_resolved" || die "Start commit '$from_commit' must be an ancestor of end commit '$to_commit' on $REMOTE_NAME/$BRANCH."

  while IFS= read -r commit; do
    [ -n "$commit" ] || continue
    if commit_is_available_for_import "$commit"; then
      selected_commits+=("$commit")
    fi
  done < <(git rev-list --reverse "${from_resolved}^..${to_resolved}")

  [ "${#selected_commits[@]}" -gt 0 ] || die "No importable commits were found in the requested range."
  run_import_commits "${selected_commits[@]}"
}

run_continue() {
  ensure_git_repo
  ensure_current_branch_with_state_fallback
  ensure_upstream_remote
  load_state || die "No saved cherry-pick state was found."
  [ "$CURRENT_BRANCH" = "${STATE_BRANCH:-$CURRENT_BRANCH}" ] || die "Saved cherry-pick state belongs to branch '$STATE_BRANCH'. Check out that branch first."
  has_cherry_pick_in_progress || die "There is no in-progress cherry-pick to continue."

  say "Continuing cherry-pick..."
  GIT_EDITOR=: git cherry-pick --continue

  if has_cherry_pick_in_progress; then
    STATE_STATUS="cherry_pick_conflict"
    update_state
    say "Cherry-pick hit another conflict."
    say "Current conflicting upstream commit: $(git show --no-patch --format='%h %s' CHERRY_PICK_HEAD)"
    say "Resolve conflicts and run: ./.import-upstream-cherry-pick.sh continue"
    say "Or return safely with: ./.import-upstream-cherry-pick.sh rollback"
    exit 1
  fi

  STASH_HASH="${STATE_STASH_HASH:-}"
  STASH_REF=""
  BACKUP_BRANCH="${STATE_BACKUP_BRANCH:-}"
  restore_stash 1
  say "Cherry-pick import completed."
}

run_restore_saved_stash() {
  ensure_git_repo
  ensure_no_git_operation_in_progress
  ensure_current_branch_with_state_fallback
  load_state || die "No saved import state was found."
  [ "$CURRENT_BRANCH" = "${STATE_BRANCH:-$CURRENT_BRANCH}" ] || die "Saved import state belongs to branch '$STATE_BRANCH'. Check out that branch first."
  [ -n "${STATE_STASH_HASH:-}" ] || die "There is no saved pre-import stash to restore."
  if worktree_has_local_changes; then
    die "The working tree is not clean. Commit, stash, or discard current changes before restoring the saved stash."
  fi

  prepare_safety_before_real_action "before-restore-stash"
  STASH_HASH="$STATE_STASH_HASH"
  STASH_REF=""
  restore_stash 1
  say "Saved pre-import local changes restored."
}

run_rollback() {
  local explicit_target="${1:-}" reset_target="" reset_snapshot=""
  local original_restore_stash_hash="" original_restore_start_head=""
  local rollback_safety_snapshot="" rollback_safety_backup="" rollback_safety_start_head=""
  local preserve_stash_hash=""

  ensure_git_repo
  ensure_no_noncherrypick_git_operation_in_progress
  load_state || true
  ensure_current_branch_with_state_fallback
  ensure_upstream_remote

  if [ -n "$explicit_target" ] && is_valid_snapshot_dir "$explicit_target"; then
    reset_snapshot="$explicit_target"
  elif [ -n "$explicit_target" ]; then
    git show-ref --verify --quiet "refs/heads/$explicit_target" || die "Backup branch '$explicit_target' does not exist."
    reset_target="$explicit_target"
  else
    if [ -n "${STATE_BRANCH:-}" ] && [ "$CURRENT_BRANCH" != "$STATE_BRANCH" ]; then
      die "Saved import state belongs to branch '$STATE_BRANCH'. Check out that branch or pass an explicit backup branch or snapshot."
    fi
    if [ -n "${STATE_SNAPSHOT_PATH:-}" ] && is_valid_snapshot_dir "$STATE_SNAPSHOT_PATH"; then
      reset_snapshot="$STATE_SNAPSHOT_PATH"
    fi
    if [ -n "${STATE_START_HEAD:-}" ]; then
      reset_target="$STATE_START_HEAD"
    elif [ -n "${STATE_BACKUP_BRANCH:-}" ]; then
      reset_target="$STATE_BACKUP_BRANCH"
    else
      reset_target="$(find_latest_backup_branch_for_current || true)"
    fi
  fi

  if [ -z "$reset_target" ] && [ -n "$reset_snapshot" ]; then
    reset_target="$(snapshot_metadata_value "$reset_snapshot" head || true)"
  fi

  [ -n "$reset_target" ] || [ -n "$reset_snapshot" ] || die "No saved rollback target was found."

  original_restore_stash_hash="${STATE_STASH_HASH:-}"
  original_restore_start_head="${STATE_START_HEAD:-}"
  if [ -z "$original_restore_start_head" ] && [ -n "$reset_snapshot" ]; then
    original_restore_start_head="$(snapshot_metadata_value "$reset_snapshot" head || true)"
  fi

  if has_cherry_pick_in_progress; then
    say "Aborting in-progress cherry-pick before rollback..."
    git cherry-pick --abort || die "Failed to abort the in-progress cherry-pick."
  fi

  if [ -n "$(git diff --name-only --diff-filter=U)" ]; then
    die "There are still unmerged files in the working tree. Resolve or discard them before rollback."
  fi

  prepare_safety_before_real_action "before-rollback"
  rollback_safety_backup="$BACKUP_BRANCH"
  rollback_safety_snapshot="$STATE_SNAPSHOT_PATH"
  rollback_safety_start_head="$STATE_START_HEAD"

  preserve_stash_hash=""
  if worktree_has_local_changes; then
    stash_current_changes "pre-rollback-preserve"
    preserve_stash_hash="$STASH_HASH"
  fi

  if [ -n "$reset_target" ]; then
    say "Resetting '$CURRENT_BRANCH' to '$reset_target'..."
    git reset --hard "$reset_target" >/dev/null
  fi

  if [ -n "$original_restore_stash_hash" ]; then
    STASH_HASH="$original_restore_stash_hash"
    STASH_REF=""
    BACKUP_BRANCH="$rollback_safety_backup"
    say "Restoring original pre-import local changes..."
    restore_stash 1
  fi

  if [ -n "$reset_snapshot" ]; then
    say "Restoring exact worktree snapshot from: $reset_snapshot"
    restore_worktree_snapshot "$reset_snapshot"
  fi

  STATE_BRANCH="$CURRENT_BRANCH"
  STATE_BACKUP_BRANCH="$rollback_safety_backup"
  STATE_STASH_HASH="$preserve_stash_hash"
  STATE_REMOTE_NAME="$REMOTE_NAME"
  STATE_REMOTE_URL="$REMOTE_URL"
  STATE_TARGET_BRANCH="$BRANCH"
  STATE_START_HEAD="$rollback_safety_start_head"
  STATE_SNAPSHOT_PATH="$rollback_safety_snapshot"
  STATE_IMPORT_COMMITS=""
  STATE_STATUS="rollback_completed"
  write_state

  say "Rollback completed."
  say "Restored starting commit: ${original_restore_start_head:-$reset_target}"
  if [ -n "$reset_snapshot" ]; then
    say "Restored snapshot: $reset_snapshot"
  fi
  say "Safety backup of the pre-rollback state kept at: $rollback_safety_backup"
  say "Safety snapshot of the pre-rollback state kept at: $rollback_safety_snapshot"
  if [ -n "$preserve_stash_hash" ]; then
    say "Safety stash of the pre-rollback local changes kept at: $preserve_stash_hash"
  fi
}

run_restore_snapshot() {
  local target_snapshot="${1:-}" restore_safety_backup="" restore_safety_snapshot=""
  local restore_safety_start_head="" restore_start_head="" preserve_stash_hash=""

  ensure_git_repo
  ensure_no_noncherrypick_git_operation_in_progress
  load_state || die "No saved import state was found."
  ensure_current_branch_with_state_fallback
  ensure_upstream_remote

  if [ -z "$target_snapshot" ]; then
    target_snapshot="${STATE_SNAPSHOT_PATH:-}"
  fi
  is_valid_snapshot_dir "$target_snapshot" || die "No valid saved snapshot was found to restore."
  restore_start_head="${STATE_START_HEAD:-}"
  if [ -z "$restore_start_head" ]; then
    restore_start_head="$(snapshot_metadata_value "$target_snapshot" head || true)"
  fi
  [ -n "$restore_start_head" ] || die "No saved starting commit was found for the snapshot restore."

  if has_cherry_pick_in_progress; then
    say "Aborting in-progress cherry-pick before snapshot restore..."
    git cherry-pick --abort || die "Failed to abort the in-progress cherry-pick."
  fi

  if [ -n "$(git diff --name-only --diff-filter=U)" ]; then
    die "There are still unmerged files in the working tree. Resolve or discard them before restoring the snapshot."
  fi

  prepare_safety_before_real_action "before-restore-snapshot"
  restore_safety_backup="$BACKUP_BRANCH"
  restore_safety_snapshot="$STATE_SNAPSHOT_PATH"
  restore_safety_start_head="$STATE_START_HEAD"

  if worktree_has_local_changes; then
    stash_current_changes "pre-restore-snapshot-preserve"
    preserve_stash_hash="$STASH_HASH"
  fi

  say "Resetting '$CURRENT_BRANCH' to '$restore_start_head' before snapshot restore..."
  git reset --hard "$restore_start_head" >/dev/null

  if [ -n "${STATE_STASH_HASH:-}" ]; then
    STASH_HASH="$STATE_STASH_HASH"
    STASH_REF=""
    BACKUP_BRANCH="$restore_safety_backup"
    say "Restoring saved pre-import local changes..."
    restore_stash 1
  fi

  say "Restoring exact worktree snapshot from: $target_snapshot"
  restore_worktree_snapshot "$target_snapshot"

  STATE_BRANCH="$CURRENT_BRANCH"
  STATE_BACKUP_BRANCH="$restore_safety_backup"
  STATE_STASH_HASH="$preserve_stash_hash"
  STATE_REMOTE_NAME="$REMOTE_NAME"
  STATE_REMOTE_URL="$REMOTE_URL"
  STATE_TARGET_BRANCH="$BRANCH"
  STATE_START_HEAD="$restore_safety_start_head"
  STATE_SNAPSHOT_PATH="$restore_safety_snapshot"
  STATE_IMPORT_COMMITS=""
  STATE_STATUS="snapshot_restored"
  write_state

  say "Snapshot restore completed."
  say "Safety backup of the pre-restore state kept at: $restore_safety_backup"
  say "Safety snapshot of the pre-restore state kept at: $restore_safety_snapshot"
  if [ -n "$preserve_stash_hash" ]; then
    say "Safety stash of the pre-restore local changes kept at: $preserve_stash_hash"
  fi
}

show_menu() {
  local choice import_commits range_from range_to snapshot_target rollback_target

  while true; do
    cat <<EOF
Available actions:
  1) Fetch upstream only
  2) List available upstream commits
  3) Import selected commit(s)
  4) Import an inclusive commit range
  5) Continue in-progress cherry-pick
  6) Rollback last import
  7) Restore saved stash
  8) Restore exact snapshot
  9) Create backup branch only
  10) Show repository status
  11) Help
  0) Exit
EOF
    printf 'Choose an action: '
    read -r choice

    case "$choice" in
      1) run_fetch_only ;;
      2) run_list ;;
      3)
        printf 'Enter upstream commit hash(es), separated by spaces: '
        read -r -a import_commits
        run_import_commits "${import_commits[@]}"
        return
        ;;
      4)
        printf 'Enter the start commit hash: '
        read -r range_from
        printf 'Enter the end commit hash: '
        read -r range_to
        run_import_range "$range_from" "$range_to"
        return
        ;;
      5)
        run_continue
        return
        ;;
      6)
        printf 'Optional backup branch or snapshot path (leave empty for saved state): '
        read -r rollback_target
        run_rollback "$rollback_target"
        return
        ;;
      7)
        run_restore_saved_stash
        return
        ;;
      8)
        printf 'Optional snapshot path (leave empty for saved state): '
        read -r snapshot_target
        run_restore_snapshot "$snapshot_target"
        return
        ;;
      9) run_backup_only ;;
      10) run_status ;;
      11) show_usage ;;
      0) exit 0 ;;
      *) say "Invalid menu selection: $choice" ;;
    esac
  done
}

main() {
  case "${1:-menu}" in
    fetch) run_fetch_only ;;
    list) run_list ;;
    import)
      shift
      run_import_commits "$@"
      ;;
    import-range) run_import_range "${2:-}" "${3:-}" ;;
    continue) run_continue ;;
    backup) run_backup_only ;;
    rollback) run_rollback "${2:-}" ;;
    restore-stash) run_restore_saved_stash ;;
    restore-snapshot) run_restore_snapshot "${2:-}" ;;
    status) run_status ;;
    help|-h|--help) show_usage ;;
    menu)
      if [ -t 0 ] && [ -t 1 ]; then
        show_menu
      else
        show_usage
      fi
      ;;
    *)
      die "Unknown command: $1. Run with 'help' to see available options."
      ;;
  esac
}

main "$@"