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

RCLONE_CONF="${HOME}/.config/rclone/rclone.conf"
SSH_DIR="${HOME}/.ssh"

if [ "$(id -u)" -eq 0 ]; then
    echo -e "\033[0;31m[WARNING]\033[0m Running as root — keys will be written to /root/.ssh. Run as armbian instead."
    exit 1
fi
mkdir -p "$SSH_DIR"
touch "$SSH_DIR/known_hosts"
chmod 700 "$SSH_DIR"
chmod 600 "$SSH_DIR/known_hosts"

check_deps ssh ssh-keygen ssh-keyscan ssh-copy-id rclone

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

retry_or_abort() {
    local test_cmd="$1"
    local fail_msg="$2"
    while true; do
        if eval "$test_cmd" 2>/dev/null; then
            return 0
        fi
        echo ""
        warn "$fail_msg"
        echo -e "  ${YELLOW}[R]${NC} Retry   ${RED}[A]${NC} Abort"
        read -rp "Choice: " CHOICE
        case "${CHOICE,,}" in
            r) info "Retrying..." ;;
            a) die "Aborted." ;;
            *) warn "Press R to retry or A to abort." ;;
        esac
    done
}

info()   { echo -e "${GREEN}[+]${NC} $*"; }
check_deps() {
    local missing=()
    for cmd in "$@"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done
    if [ ${#missing[@]} -gt 0 ]; then
        die "Missing required dependencies: ${missing[*]}"
    fi
}


warn()   { echo -e "${YELLOW}[!]${NC} $*"; }
die()    { echo -e "${RED}[ERROR]${NC} $*" >&2; exit 1; }
header() { echo -e "\n${CYAN}══════════════════════════════════════════${NC}"; echo -e "${CYAN}  $*${NC}"; echo -e "${CYAN}══════════════════════════════════════════${NC}"; }

header "TinyBoard Hub — Onboard New Spoke"

read -rp "Spoke name (e.g. rocky): " SPOKE_NAME
[ -n "$SPOKE_NAME" ] || die "Spoke name cannot be empty"

read -rp "Tunnel port for $SPOKE_NAME: " TUNNEL_PORT
[[ "$TUNNEL_PORT" =~ ^[0-9]+$ ]] || die "Invalid port"

KEY_NAME="armbian-${SPOKE_NAME}-$(date +%Y%m)"
KEY_PATH="$SSH_DIR/$KEY_NAME"

command -v rclone >/dev/null || die "rclone is not installed"
mkdir -p "$(dirname "$RCLONE_CONF")"

header "Checking Tunnel"
info "Scanning spoke host key..."
KEYSCAN=$(ssh-keyscan -p "$TUNNEL_PORT" -H localhost 2>/dev/null)
[ -n "$KEYSCAN" ] || die "Spoke not reachable on port $TUNNEL_PORT — is the tunnel up?"
echo "$KEYSCAN" >> "$SSH_DIR/known_hosts"

info "Verifying spoke is reachable on port $TUNNEL_PORT..."
retry_or_abort \
    "ssh -o BatchMode=yes -o ConnectTimeout=10 -p \"$TUNNEL_PORT\" armbian@localhost exit" \
    "Spoke not reachable on port $TUNNEL_PORT. Make sure the tunnel is up."

header "Generating Hub SSH Key"
if [ -f "$KEY_PATH" ]; then
    warn "Key $KEY_PATH already exists, skipping generation."
else
    ssh-keygen -t ed25519 -f "$KEY_PATH" -N ""
    info "Key generated: $KEY_PATH"
fi

header "Copying Hub Key to Spoke"
info "Running ssh-copy-id to armbian@localhost:$TUNNEL_PORT..."
info "(You will be prompted for the armbian password on the spoke)"
ssh-copy-id -i "$KEY_PATH.pub" -p "$TUNNEL_PORT" armbian@localhost
info "Key copied."

header "Testing Hub -> Spoke Key Auth"
retry_or_abort \
    "ssh -i \"$KEY_PATH\" -o BatchMode=yes -o ConnectTimeout=10 -p \"$TUNNEL_PORT\" armbian@localhost exit" \
    "Key auth failed. Check authorized_keys on the spoke."
info "Key auth to spoke successful."

header "Adding rclone Remote"
if grep -q "\[${SPOKE_NAME}-remote\]" "$RCLONE_CONF" 2>/dev/null; then
    warn "Remote [${SPOKE_NAME}-remote] already exists in $RCLONE_CONF, skipping."
else
    cat >> "$RCLONE_CONF" <<EOF

[${SPOKE_NAME}-remote]
type = sftp
host = localhost
port = $TUNNEL_PORT
key_file = $KEY_PATH
shell_type = unix
md5sum_command = md5sum
sha1sum_command = sha1sum
EOF
    info "Remote [${SPOKE_NAME}-remote] added to $RCLONE_CONF."
fi

header "Testing rclone Connection"
if rclone lsd "${SPOKE_NAME}-remote:" --config "$RCLONE_CONF" 2>/dev/null; then
    info "rclone connection to $SPOKE_NAME successful."
else
    warn "rclone test failed. Check the remote config in $RCLONE_CONF."
fi

header "Onboarding Complete"
echo -e "  Spoke:       ${GREEN}$SPOKE_NAME${NC}"
echo -e "  Port:        ${GREEN}$TUNNEL_PORT${NC}"
echo -e "  Hub key:     ${GREEN}$KEY_PATH${NC}"
echo -e "  rclone:      ${GREEN}${SPOKE_NAME}-remote${NC}"
echo ""
echo -e "${YELLOW}To mount this spoke:${NC}"
echo "  RCLONE_REMOTE=${SPOKE_NAME}-remote hubspoke-helper.sh hub start"
echo ""