#!/bin/bash
# Keyboard-driven SSH hub (fzf). Launched by Mod+S in a floating kitty (class ssh-menu).
#   Enter ssh   Alt+T ssh+tmux   Alt+F sftp   Alt+Y copy addr   Alt+K push key
#   Alt+N new host   Alt+E edit   Alt+G group   Alt+R toggle PROD   Tab mark   Alt+A/Z all/none
# Host metadata: "# group: X" / "# env: prod" comments above Host lines in ~/.ssh/config.
S=$(readlink -f "$0")
H="python3 -I $(dirname "$S")/ssh_hosts.py"
state=${XDG_STATE_HOME:-$HOME/.local/state}/ssh-menu
mkdir -p "$state"; touch "$state/history"

c_ok=$'\e[32m'; c_err=$'\e[31m'; c_dim=$'\e[2m'; c_hi=$'\e[38;2;233;30;99m'; c_off=$'\e[0m'
say()  { printf '%s\n' "$*"; }
ask()  { local v; read -e -r -p "$c_hi$1$c_off ${2:+[$2] }" v; printf '%s' "${v:-$2}"; }
yes()  { local v; read -r -p "$c_hi$1$c_off [Y/n] " v; [[ ! $v =~ ^[nNнН] ]]; }
pause(){ read -r -s -n1 -p "${c_dim}press any key${c_off}"; echo; }
fz()   { fzf --layout=reverse --height=45% --border=rounded --prompt="$1" --pointer='▌' \
              --color='bg+:#261d1e,fg+:#ffffff,hl:#e91e63,hl+:#ff6398,pointer:#e91e63,prompt:#e91e63,border:#e91e63,info:#8e7a7d' "${@:2}"; }

# ---------- connect ----------
run_ssh() {  # run_ssh <target> [remote command...]
    local t=$1; shift
    echo "$t" >> "$state/history"
    # become a normal tiled window once we are a session, not a menu
    [ "$SSH_MENU" = 1 ] && hyprctl dispatch 'hl.dsp.window.float({ action = "disable" })' >/dev/null 2>&1
    local env; env=$($H env "$t")
    if [ "$env" = prod ]; then printf '\e]2;⚠ PROD ssh: %s\a\e]11;#220a0f\a' "$t"
    else printf '\e]2;ssh: %s\a' "$t"; fi
    ssh "$t" "$@"; local rc=$?
    printf '\e]111\a'
    if [ $rc -ne 0 ]; then
        printf '\n%sssh %s exited with %s%s  (press Enter to close)' "$c_err" "$t" "$rc" "$c_off"; read -r _
    fi
    exit $rc
}
cmd_connect() { local t=${1:-$2}; [ -n "$t" ] && run_ssh "$t"; }
cmd_tmux()    { run_ssh "$1" -t 'command -v tmux >/dev/null && exec tmux new-session -A -s main || exec $SHELL -l'; }

# ---------- quick actions ----------
cmd_sftp() {
    local a jump url
    read -r a jump < <($H addr "$1")
    # jump hosts: let gvfs go through ssh config by alias; otherwise explicit user@host:port
    if [ -n "$jump" ]; then url="sftp://$1/"; else url="sftp://$a/"; fi
    notify-send -a ssh-menu "SFTP" "opening $a …"
    # detached, so it survives the menu window closing
    setsid -f nautilus --new-window "$url" >/dev/null 2>&1
}
cmd_copy() { read -r a _ < <($H addr "$1"); printf %s "$a" | wl-copy; notify-send -a ssh-menu "ssh" "copied $a"; }

# ---------- keys ----------
pick_key() {
    local pubs sel name
    mapfile -t pubs < <(ls "$HOME"/.ssh/*.pub 2>/dev/null)
    sel=$( { printf '%s\n' "${pubs[@]}"; echo "+ generate new ed25519 key"; } | grep -v '^$' | fz 'key ❯ ' --height=30%) || return 1
    if [[ $sel == "+ generate"* ]]; then
        name=$(ask "key file name" "id_ed25519_$(date +%y%m)")
        [ -e "$HOME/.ssh/$name" ] && { say "${c_err}$name exists${c_off}" >&2; return 1; }
        ssh-keygen -t ed25519 -f "$HOME/.ssh/$name" -C "$USER@$(hostname -s)" >&2 || return 1
        sel="$HOME/.ssh/$name.pub"
    fi
    printf '%s' "$sel"
}
cmd_keys() {
    [ $# -eq 0 ] && return
    local pub ok=() bad=() h
    pub=$(pick_key) || return
    say "${c_hi}key:${c_off} $pub  ${c_hi}→${c_off} $*"; echo
    for h in "$@"; do
        say "${c_hi}── $h${c_off}"
        if ssh-copy-id -i "$pub" "$h"; then ok+=("$h"); else bad+=("$h"); fi
    done
    echo; [ ${#ok[@]} -gt 0 ]  && say "${c_ok}✓ ${ok[*]}${c_off}"
    [ ${#bad[@]} -gt 0 ] && say "${c_err}✗ ${bad[*]}${c_off}"
    pause
}

# ---------- new host / edit / group ----------
cmd_new() {
    clear; say "${c_hi}── new SSH host ──${c_off}  (Enter = default)"; echo
    local alias hostname user port jump group env key
    while :; do
        alias=$(ask "alias (name in the menu)"); [ -z "$alias" ] && return
        $H has "$alias" && { say "${c_err}already exists${c_off}"; continue; }
        [[ $alias =~ [[:space:]*?!] ]] && { say "${c_err}no spaces or wildcards${c_off}"; continue; }
        break
    done
    hostname=$(ask "HostName / IP" "$alias")
    user=$(ask "User" "root")
    port=$(ask "Port" "22")
    jump=$(ask "ProxyJump (empty = none)")
    group=$(ask "group (empty = none)")
    yes "mark as PROD (red window)?" && env=prod || env=
    key=$(ask "IdentityFile (empty = default)")
    echo; say "${c_dim}Host $alias → $user@$hostname:$port ${jump:+via $jump }${group:+[$group] }${env:+PROD}${c_off}"
    yes "save to ~/.ssh/config?" || return
    $H add "$alias" "$hostname" "$user" "$port" "$jump" "$group" "$env" "$key" || { pause; return; }
    say "${c_ok}saved${c_off} (backup in $state/backups)"
    yes "push your public key to it now?" && cmd_keys "$alias"
    echo "$alias" > "$state/lastnew"
}
cmd_edit() {
    local ed=micro; command -v micro >/dev/null || ed=${EDITOR:-nano}
    local l; l=$($H line "$1")
    if [ "$ed" = micro ]; then micro "+$l" ~/.ssh/config; else "$ed" "+$l" ~/.ssh/config; fi
}
cmd_group() {
    [ $# -eq 0 ] && return
    local out rc g h
    out=$($H groups | fz "group ❯ " --height=40% --print-query \
        --header "pick a group or type a new name (empty = clear)"); rc=$?
    [ $rc -gt 1 ] && return          # Esc
    mapfile -t o <<< "$out"          # line 1 = query, line 2 = selection (if any)
    g=${o[1]:-${o[0]}}
    for h in "$@"; do $H set "$h" group "$g"; done
}
cmd_envtoggle() { for h in "$@"; do $H toggle-prod "$h"; done; }

# ---------- several hosts ----------
confirm_prod() {  # confirm_prod host... ; asks (default: no) only when PROD hosts are involved
    local prod=() h v
    for h in "$@"; do [ "$($H env "$h")" = prod ] && prod+=("$h"); done
    [ ${#prod[@]} -eq 0 ] && return 0
    say "${c_err}${#prod[@]} of $# hosts are PROD:${c_off} ${prod[*]}"
    read -r -p "${c_hi}continue?${c_off} [y/N] " v
    [[ $v =~ ^[yYдД] ]]
}
cmd_broadcast() {  # one tmux pane per host, keyboard input synced (Alt+S toggles)
    [ $# -eq 0 ] && return
    confirm_prod "$@" || return
    local conf; conf="$(dirname "$S")/tmux-broadcast.conf"
    local T=(tmux -L sshmenu -f "$conf") sess="bc-$(date +%H%M%S)" h first=1
    for h in "$@"; do
        echo "$h" >> "$state/history"
        if [ $first = 1 ]; then "${T[@]}" new-session -d -s "$sess" -x 220 -y 60 "ssh $(printf %q "$h")"; first=0
        else "${T[@]}" split-window -t "$sess" "ssh $(printf %q "$h")"; "${T[@]}" select-layout -t "$sess" tiled; fi
        "${T[@]}" select-pane -t "$sess" -T "$h"
    done
    "${T[@]}" setw -t "$sess" synchronize-panes on
    setsid -f kitty --title "broadcast: $# hosts" "${T[@]}" attach -t "$sess" >/dev/null 2>&1
}
cmd_run() {  # run one command on all marked hosts, identical answers merged
    [ $# -eq 0 ] && return
    local cmd v
    set -o history; HISTFILE="$state/cmd_history"; history -r
    say "${c_dim}run on: $*${c_off}"
    read -e -r -p "${c_hi}\$${c_off} " cmd; [ -z "$cmd" ] && return
    history -s "$cmd"; history -w
    confirm_prod "$@" || return
    clear
    $H run "$cmd" "$@" | less -RFX
    pause
}
cmd_web() {  # open the host's web UI; first time asks for the URL and stores it as "# url:"
    local url; url=$($H url "$1")
    if [ -z "$url" ]; then
        url=$(ask "web UI for $1" "$($H url-default "$1")"); [ -z "$url" ] && return
        $H set "$1" url "$url"
    fi
    # Chrome (Flatpak) if installed, otherwise the default browser
    if flatpak info com.google.Chrome >/dev/null 2>&1; then setsid -f flatpak run com.google.Chrome "$url" >/dev/null 2>&1
    else setsid -f xdg-open "$url" >/dev/null 2>&1; fi
}

# ---------- dispatcher ----------
case "$1" in
    connect)    shift; cmd_connect "$@"; exit ;;
    tmux)       shift; cmd_tmux "$@"; exit ;;
    sftp)       shift; cmd_sftp "$@"; exit ;;
    copy)       shift; cmd_copy "$@"; exit ;;
    keys)       shift; cmd_keys "$@"; exit ;;
    new)        cmd_new; exit ;;
    edit)       shift; cmd_edit "$@"; exit ;;
    group)      shift; cmd_group "$@"; exit ;;
    envtoggle)  shift; cmd_envtoggle "$@"; exit ;;
    broadcast)  shift; cmd_broadcast "$@"; exit ;;
    run)        shift; cmd_run "$@"; exit ;;
    web)        shift; cmd_web "$@"; exit ;;
    preview)    shift; $H preview "$1"; exit ;;
    list)       $H list; exit ;;
    live)       $H live; exit ;;
esac

# ---------- menu ----------
# hotkey table: accent key + dim description, 4 columns
key() { printf '\e[1;38;2;248;187;217m%-6s\e[0m\e[38;2;150;135;140m%-15s\e[0m' "$1" "$2"; }
hdr=$(
    key Enter connect;    key Alt+T tmux;        key Alt+F sftp;         key Alt+W web;       echo
    key Alt+Y 'copy addr'; key Alt+N 'new host'; key Alt+E 'edit config'; key Alt+R 'toggle PROD'; echo
    key Tab mark;         key Alt+A 'mark all';  key Alt+Z unmark;       key Alt+G group;     echo
    key Alt+B broadcast;  key Alt+X 'run cmd';   key Alt+K 'push key'
)

$H list | fzf --ansi --multi --layout=reverse --height=100% --tiebreak=index \
    --style=full:rounded --padding=1,2 --highlight-line \
    --prompt='  ' --pointer='▌' --gutter=' ' --marker='✓' --info=inline-right --no-scrollbar \
    --input-label=' ssh ' --list-label=' hosts ' --header-label=' keys ' --preview-label=' details ' \
    --ghost='type to filter, or user@host to connect to anything' \
    --header "$hdr" \
    --color='bg+:#2a2024,fg+:#ffffff,hl:#e91e63,hl+:#ff6398,pointer:#e91e63,prompt:#e91e63,info:#8e7a7d,marker:#f8bbd9,spinner:#e91e63,ghost:#6b5a5e,gutter:-1' \
    --color='border:#3a2a2e,list-border:#3a2a2e,header-border:#3a2a2e,preview-border:#3a2a2e,input-border:#e91e63' \
    --color='label:#e91e63,input-label:#e91e63,list-label:#8e7a7d,header-label:#8e7a7d,preview-label:#8e7a7d,header:-1' \
    --preview "$S preview {2}" --preview-window='right:32%:wrap' \
    --bind "start:reload($S live)" \
    --bind "enter:become($S connect {2} {q})" \
    --bind "alt-t:become($S tmux {2})" \
    --bind "alt-f:execute-silent($S sftp {2})+abort" \
    --bind "alt-y:execute-silent($S copy {2})+abort" \
    --bind "alt-k:execute($S keys {+2})+reload($S list)" \
    --bind "alt-n:execute($S new)+reload($S list)+transform-query(cat $state/lastnew 2>/dev/null; rm -f $state/lastnew)" \
    --bind "alt-e:execute($S edit {2})+reload($S live)" \
    --bind "alt-g:execute($S group {+2})+reload($S list)" \
    --bind "alt-r:execute-silent($S envtoggle {+2})+reload($S list)" \
    --bind "alt-b:execute($S broadcast {+2})+abort" \
    --bind "alt-x:execute($S run {+2})" \
    --bind "alt-w:execute($S web {2})+reload($S list)+abort" \
    --bind "alt-a:select-all,alt-z:deselect-all"