#!/usr/bin/env bash
#
# JetBrains IDE plugin installer for macOS, Linux (Ubuntu Desktop) and WSL2.
#
# Detects every locally installed JetBrains IDE (Toolbox + standalone) and,
# on Linux, any JetBrains Gateway remote-dev-server.sh distributions. Lets
# you multi-select IDE(s) and plugin(s) interactively, then installs each
# plugin into each selected IDE.
#
# Picker UX:
#   - Uses fzf with --multi when fzf is on PATH (TAB to toggle, Enter to confirm).
#   - Otherwise prints a numbered list and accepts space-separated indexes
#     (e.g. "1 3 5"), exclusions (e.g. "!1 !4" to select all except 1 and 4), 
#     or "a" for all.
#
# Usage:
#   ./install-jetbrains.sh        # interactive
#   ./install-jetbrains.sh --help

set -euo pipefail
IFS=$'\n\t'

# ------------------------------------------------------------------
# Canonical marketplace plugin list (xmlId | display label).
# Both columns separated by a literal pipe; the label is for the picker.
# ------------------------------------------------------------------
readonly PLUGIN_CATALOG=(
  "com.intellij.ml.llm|JetBrains AI Assistant"
  "izhangzhihao.rainbow.brackets|Rainbow Brackets"
  "IdeaVIM|IdeaVim"
  "org.sonarlint.idea|SonarQube for IDE (SonarLint)"
  "Key Promoter X|Key Promoter X"
  "net.ashald.envfile|EnvFile"
  "org.intellij.qodana|Qodana"
)

# Known JetBrains IDE launcher basenames (Toolbox shims + Linux .sh names).
readonly KNOWN_IDE_NAMES=(
  idea idea.sh
  pycharm pycharm.sh
  webstorm webstorm.sh
  goland goland.sh
  rubymine rubymine.sh
  clion clion.sh
  phpstorm phpstorm.sh
  datagrip datagrip.sh
  rustrover rustrover.sh
  rider rider.sh
  studio studio.sh
  fleet
)

usage() {
  cat <<EOF
Usage: $(basename "$0")

Interactive JetBrains IDE plugin installer for macOS, Linux and WSL2.

The script:
  1. Detects your OS (macOS / Linux / WSL2)
  2. Scans for installed JetBrains IDEs (Toolbox + standalone) and any
     local remote-dev-server.sh distributions
  3. Lets you multi-select IDE(s) and plugin(s)
  4. Runs each launcher's "installPlugins" command for every selected plugin

Use \`brew install fzf\` (or your distro's package manager) for a nicer picker.
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    -h|--help) usage; exit 0 ;;
    *)         echo "Error: unknown argument '$1'." >&2; usage >&2; exit 2 ;;
  esac
done

# ------------------------------------------------------------------
# OS detection
# ------------------------------------------------------------------
OS_KIND=""
case "$(uname -s)" in
  Darwin) OS_KIND=macos ;;
  Linux)
    if grep -qi microsoft /proc/version 2>/dev/null; then
      OS_KIND=wsl2
    else
      OS_KIND=linux
    fi
    ;;
  *)
    echo "Error: unsupported OS '$(uname -s)'. Use install-jetbrains.ps1 on Windows." >&2
    exit 1
    ;;
esac

# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------

# is_known_ide_name <basename>  -> 0 if it matches a known JetBrains launcher.
is_known_ide_name() {
  local n="$1"
  local k
  for k in "${KNOWN_IDE_NAMES[@]}"; do
    [[ "$n" == "$k" ]] && return 0
  done
  return 1
}

# pretty_from_path /path/to/idea.sh  -> "IntelliJ IDEA (~/.local/share/.../idea.sh)"
# Builds a short, readable label for the picker.
pretty_label() {
  local path="$1" base
  base="$(basename "$path")"
  base="${base%.sh}"
  base="${base%.cmd}"
  base="${base%.exe}"
  case "$base" in
    idea|idea64)        printf 'IntelliJ IDEA'   ;;
    pycharm|pycharm64)  printf 'PyCharm'         ;;
    webstorm|webstorm64) printf 'WebStorm'       ;;
    goland|goland64)    printf 'GoLand'          ;;
    rubymine|rubymine64) printf 'RubyMine'       ;;
    clion|clion64)      printf 'CLion'           ;;
    phpstorm|phpstorm64) printf 'PhpStorm'       ;;
    datagrip|datagrip64) printf 'DataGrip'       ;;
    rustrover|rustrover64) printf 'RustRover'    ;;
    rider|rider64)      printf 'Rider'           ;;
    studio|studio64)    printf 'Android Studio'  ;;
    fleet)              printf 'Fleet'           ;;
    *)                  printf '%s' "$base"      ;;
  esac
  printf '  (%s)' "$path"
}

# Translate a Linux/WSL home into the Windows-side user dir under /mnt/c.
win_user_dir() {
  local winuser
  winuser="$(cmd.exe /c 'echo %USERNAME%' 2>/dev/null | tr -d '\r\n' || true)"
  [[ -z "$winuser" ]] && winuser="$USER"
  printf '/mnt/c/Users/%s' "$winuser"
}

# ------------------------------------------------------------------
# IDE candidate scan
# Each candidate is recorded as a single line:
#     <type>|<launcher>|<label>
# where <type> is one of:
#     ide        -> run "<launcher> installPlugins <id>..."
#     wincmd     -> run "cmd.exe /c <launcher> installPlugins <id>..."
#     remotedev  -> remote-dev-server.sh installPlugins ...
# ------------------------------------------------------------------
declare -a CANDIDATES=()

add_candidate() {
  local type="$1" launcher="$2"
  CANDIDATES+=("${type}|${launcher}|$(pretty_label "$launcher")")
}

scan_macos() {
  local d="$HOME/Library/Application Support/JetBrains/Toolbox/scripts"
  if [[ -d "$d" ]]; then
    local f
    while IFS= read -r -d '' f; do
      local base
      base="$(basename "$f")"
      is_known_ide_name "$base" || continue
      [[ -x "$f" ]] && add_candidate ide "$f"
    done < <(find "$d" -maxdepth 1 -type f -print0 2>/dev/null)
  fi

  local app bin
  for app in /Applications/*.app; do
    [[ -d "$app/Contents/MacOS" ]] || continue
    [[ "$(basename "$app")" =~ ^(IntelliJ|PyCharm|WebStorm|GoLand|RubyMine|CLion|PhpStorm|DataGrip|RustRover|Rider|Android\ Studio|Fleet) ]] || continue
    for bin in "$app"/Contents/MacOS/*; do
      [[ -x "$bin" ]] || continue
      local b
      b="$(basename "$bin")"
      is_known_ide_name "$b" || continue
      add_candidate ide "$bin"
    done
  done
}

scan_linux() {
  local apps="$HOME/.local/share/JetBrains/Toolbox/apps"
  if [[ -d "$apps" ]]; then
    local f
    while IFS= read -r -d '' f; do
      local b
      b="$(basename "$f")"
      is_known_ide_name "$b" || continue
      [[ -x "$f" ]] && add_candidate ide "$f"
    done < <(find "$apps" -maxdepth 3 -type f -name '*.sh' -print0 2>/dev/null)
  fi

  # /opt/<ide>/bin/<ide>.sh — common for .deb / snap / tarball installs.
  local opt
  for opt in /opt/*/bin/*.sh; do
    [[ -x "$opt" ]] || continue
    local b
    b="$(basename "$opt")"
    is_known_ide_name "$b" || continue
    add_candidate ide "$opt"
  done

  # Remote-dev-server distributions (Orbstack/Gateway).
  local rd
  while IFS= read -r -d '' rd; do
    add_candidate remotedev "$rd"
  done < <(find "$HOME/.cache/JetBrains/RemoteDev/dist" -maxdepth 3 -type f -name 'remote-dev-server.sh' -print0 2>/dev/null)
}

scan_wsl2() {
  local winhome
  winhome="$(win_user_dir)"
  [[ -d "$winhome" ]] || return 0

  # Standalone installs under AppData\Local\Programs\<IDE>\bin\*.exe
  local exe b
  while IFS= read -r -d '' exe; do
    b="$(basename "$exe")"
    case "$b" in
      idea64.exe|pycharm64.exe|webstorm64.exe|goland64.exe|rubymine64.exe|\
      clion64.exe|phpstorm64.exe|datagrip64.exe|rustrover64.exe|rider64.exe|\
      studio64.exe|fleet.exe)
        add_candidate ide "$exe"
        ;;
    esac
  done < <(find "$winhome/AppData/Local/Programs" -maxdepth 4 -type f -name '*.exe' -print0 2>/dev/null)

  # Toolbox shell shims (.cmd) — invoke via cmd.exe.
  local scripts="$winhome/AppData/Local/JetBrains/Toolbox/scripts"
  if [[ -d "$scripts" ]]; then
    local f
    while IFS= read -r -d '' f; do
      b="$(basename "${f%.cmd}")"
      is_known_ide_name "$b" || continue
      add_candidate wincmd "$f"
    done < <(find "$scripts" -maxdepth 1 -type f -name '*.cmd' -print0 2>/dev/null)
  fi

  # Local Linux-side remote-dev-server.sh (still useful inside WSL2).
  local rd
  while IFS= read -r -d '' rd; do
    add_candidate remotedev "$rd"
  done < <(find "$HOME/.cache/JetBrains/RemoteDev/dist" -maxdepth 3 -type f -name 'remote-dev-server.sh' -print0 2>/dev/null)
}

case "$OS_KIND" in
  macos) scan_macos ;;
  linux) scan_linux ;;
  wsl2)  scan_wsl2  ;;
esac

if [[ ${#CANDIDATES[@]} -eq 0 ]]; then
  echo "No JetBrains IDEs detected on this system." >&2
  echo "Install one via the JetBrains Toolbox, then re-run." >&2
  exit 1
fi

# ------------------------------------------------------------------
# Multi-select picker (Handles Space separation & Exclusions via !)
# Reads newline-separated options on stdin, prints selected lines on stdout.
# ------------------------------------------------------------------
pick_multi() {
  local prompt="$1" header="$2"
  if command -v fzf >/dev/null 2>&1; then
    fzf --multi --reverse --prompt="$prompt " --header="$header"
    return
  fi

  # Fallback: numbered list + space-separated input with exclusion support.
  local -a items=()
  local line
  while IFS= read -r line; do items+=("$line"); done

  local total=${#items[@]} i
  while :; do
    {
      echo "$header"
      for (( i=0; i<total; i++ )); do
        printf '  %2d) %s\n' "$((i+1))" "${items[i]}"
      done
      printf '%s' "$prompt (e.g. 1 3, !2 to exclude 2, or a for all): "
    } >&2

    local input
    if ! IFS= read -r input </dev/tty; then
      echo "no tty available; install fzf or run interactively" >&2
      return 1
    fi

    if [[ "$input" == "a" || "$input" == "A" ]]; then
      printf '%s\n' "${items[@]}"
      return
    fi
    [[ -z "$input" ]] && { echo "  please enter at least one number, an exclusion, or 'a'." >&2; continue; }

    local -a toks=()
    IFS=' ' read -r -a toks <<<"$input"

    # Determine if this is an exclusion request (checks if any token starts with !)
    local is_exclusion=0
    local t
    for t in "${toks[@]}"; do
      if [[ "$t" =~ ^\! ]]; then
        is_exclusion=1
        break
      fi
    done

    local -a picked=()
    local bad=0

    if (( is_exclusion )); then
      # Track which indices are explicitly excluded using an associative array map
      local -A excluded_map=()
      
      for t in "${toks[@]}"; do
        [[ -z "$t" ]] && continue
        if [[ ! "$t" =~ ^\![0-9]+$ ]]; then
          echo "  Error: When using exclusions, all items must start with '!' (e.g., !1 !3)" >&2
          bad=1
          break
        fi
        
        local idx="${t#\!}"  # Strip the '!' character
        if (( idx < 1 || idx > total )); then bad=1; break; fi
        excluded_map[$((idx-1))]=1
      done
      
      if (( bad )); then
        echo "  invalid exclusion input; try again." >&2
        continue
      fi

      # Build final selection list by adding everything NOT in our exclusion map
      for (( i=0; i<total; i++ )); do
        if [[ -z "${excluded_map[$i]+abc}" ]]; then
          picked+=("${items[i]}")
        fi
      done
    else
      # Standard additive parsing (space-separated)
      for t in "${toks[@]}"; do
        t="${t// /}"
        [[ -z "$t" ]] && continue
        if [[ ! "$t" =~ ^[0-9]+$ ]]; then bad=1; break; fi
        if (( t < 1 || t > total )); then bad=1; break; fi
        picked+=("${items[t-1]}")
      done
    fi

    if (( bad )); then
      echo "  invalid input; try again." >&2
      continue
    fi
    if (( ${#picked[@]} == 0 )); then
      echo "  no valid selections; try again." >&2
      continue
    fi
    printf '%s\n' "${picked[@]}"
    return
  done
}

# ------------------------------------------------------------------
# IDE picker
# ------------------------------------------------------------------
declare -a ide_lines=()
for c in "${CANDIDATES[@]}"; do
  # human-readable picker line: just the third field
  ide_lines+=("${c#*|*|}")
done

# Map label -> (type, launcher) so we can recover after picking.
declare -A label_type label_launcher
for c in "${CANDIDATES[@]}"; do
  IFS='|' read -r t l lab <<<"$c"
  label_type[$lab]="$t"
  label_launcher[$lab]="$l"
done

echo "Detected ${#CANDIDATES[@]} JetBrains target(s)."
mapfile -t SELECTED_IDES < <(
  printf '%s\n' "${ide_lines[@]}" |
    pick_multi "IDEs>" "Pick IDE(s) — TAB to toggle (fzf) or spaced indexes"
)

if (( ${#SELECTED_IDES[@]} == 0 )); then
  echo "No IDEs selected; exiting." >&2
  exit 1
fi

# ------------------------------------------------------------------
# Plugin picker
# ------------------------------------------------------------------
declare -a plugin_lines=()
declare -A label_to_id
for entry in "${PLUGIN_CATALOG[@]}"; do
  IFS='|' read -r pid plabel <<<"$entry"
  line="$(printf '%-32s  %s' "$plabel" "$pid")"
  plugin_lines+=("$line")
  label_to_id[$line]="$pid"
done

mapfile -t SELECTED_PLUGIN_LINES < <(
  printf '%s\n' "${plugin_lines[@]}" |
    pick_multi "Plugins>" "Pick plugin(s) — TAB to toggle (fzf) or spaced indexes"
)

if (( ${#SELECTED_PLUGIN_LINES[@]} == 0 )); then
  echo "No plugins selected; exiting." >&2
  exit 1
fi

declare -a SELECTED_PLUGIN_IDS=()
for line in "${SELECTED_PLUGIN_LINES[@]}"; do
  SELECTED_PLUGIN_IDS+=("${label_to_id[$line]}")
done

# ------------------------------------------------------------------
# Running-IDE preflight
# JetBrains' `installPlugins` refuses to run while the IDE GUI is open
# ("Only one instance of IDEA can be run at a time."). We do a best-effort
# pgrep against the IDE app directory so the user can close them upfront
# instead of seeing N identical errors during the install loop.
# ------------------------------------------------------------------
ide_is_running() {
  local launcher="$1" type="$2"

  # remote-dev-server.sh spawns short-lived workers; WSL2->Windows lookups
  # would need tasklist.exe — skip in both cases.
  case "$type" in
    remotedev|wincmd) return 1 ;;
  esac

  case "$launcher" in
    */Toolbox/apps/*)
      # /.../Toolbox/apps/<ide-dir>/[ch-N/<ver>/]bin/idea.sh -> match "<ide-dir>"
      local rest="${launcher#*/Toolbox/apps/}"
      local ide_dir="${rest%%/*}"
      [[ -n "$ide_dir" ]] && pgrep -f -- "/Toolbox/apps/${ide_dir}/" >/dev/null 2>&1
      ;;
    */Applications/*.app/*)
      local app_path="${launcher%/Contents/MacOS/*}"
      pgrep -f -- "${app_path}/Contents/" >/dev/null 2>&1
      ;;
    *)
      local b
      b="$(basename "$launcher")"
      b="${b%.sh}"
      pgrep -fi -- "$b" >/dev/null 2>&1
      ;;
  esac
}

while :; do
  declare -a _running=()
  for ide_label in "${SELECTED_IDES[@]}"; do
    if ide_is_running "${label_launcher[$ide_label]}" "${label_type[$ide_label]}"; then
      _running+=("$ide_label")
    fi
  done
  (( ${#_running[@]} == 0 )) && break

  {
    echo
    echo "The following IDE(s) appear to be running — installPlugins needs them closed:"
    for r in "${_running[@]}"; do
      echo "  - $r"
    done
    echo
    printf 'Close them, then press Enter to retry (or Ctrl+C to abort): '
  } >&2
  if ! IFS= read -r _ </dev/tty; then
    echo "no tty; aborting" >&2
    exit 1
  fi
done

# ------------------------------------------------------------------
# Install loop
# ------------------------------------------------------------------
LOG_FILE="$(mktemp -t jb-install.XXXXXX)"
trap 'rm -f "$LOG_FILE"' EXIT

installed=0
failed=0
declare -a FAILED_LIST=()

echo
echo "Installing ${#SELECTED_PLUGIN_IDS[@]} plugin(s) into ${#SELECTED_IDES[@]} IDE(s)..."
echo

handle_single_instance_abort() {
  local ide_label="$1" current_pid="$2"
  shift 2
  echo
  echo "  ! '${ide_label}' is still running. Skipping remaining plugins for this IDE." >&2
  echo "  ! Close the IDE and re-run the script to finish."                           >&2
  # Mark current and remaining-for-this-IDE plugins as failed.
  local mark_from=0 p
  for p in "${SELECTED_PLUGIN_IDS[@]}"; do
    if (( mark_from )) || [[ "$p" == "$current_pid" ]]; then
      mark_from=1
      failed=$((failed + 1))
      FAILED_LIST+=("${ide_label} :: ${p}  (IDE running — close it and re-run)")
    fi
  done
}

for ide_label in "${SELECTED_IDES[@]}"; do
  itype="${label_type[$ide_label]}"
  launcher="${label_launcher[$ide_label]}"

  echo ">>> $ide_label"
  ide_aborted=0
  for pid in "${SELECTED_PLUGIN_IDS[@]}"; do
    (( ide_aborted )) && break
    printf '        %-34s ' "$pid"
    rc=0
    case "$itype" in
      ide|remotedev)
        "$launcher" installPlugins "$pid" >"$LOG_FILE" 2>&1 || rc=$?
        ;;
      wincmd)
        cmd.exe /c "$launcher" installPlugins "$pid" >"$LOG_FILE" 2>&1 || rc=$?
        ;;
    esac

    if (( rc == 0 )); then
      echo "ok"
      installed=$((installed + 1))
    elif grep -q "Only one instance" "$LOG_FILE" 2>/dev/null; then
      echo "SKIPPED"
      handle_single_instance_abort "$ide_label" "$pid"
      ide_aborted=1
    else
      echo "FAILED"
      failed=$((failed + 1))
      FAILED_LIST+=("${ide_label} :: ${pid}")
      sed 's/^/         /' "$LOG_FILE" >&2 || true
    fi
  done
  echo
done

echo "Done. Installed: ${installed}, Failed: ${failed}"

if (( failed > 0 )); then
  echo "Failed:" >&2
  for f in "${FAILED_LIST[@]}"; do
    echo "  - $f" >&2
  done
  exit 1
fi