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

HUB_HOST=""
HUB_USER="armbian"
ARMBIAN_HOME="/home/armbian"
SSH_DIR="$ARMBIAN_HOME/.ssh"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SPOKE_DIR="$SCRIPT_DIR/spoke"
COMPOSE="$SPOKE_DIR/compose.yaml"
START_PORT=11111

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} $*"; }
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}"; }

[ "$(id -u)" -eq 0 ] || die "Run as root"

read -rp "Hub hostname [oily.dad]: " HUB_HOST
HUB_HOST="${HUB_HOST:-oily.dad}"

header "TinyBoard Spoke Setup"

info "Installing packages..."
apt-get update -q
apt-get install -y -q vim autossh docker.io docker-compose-plugin git openssh-server

info "Adding armbian to docker group..."
usermod -aG docker armbian 2>/dev/null || true

info "Enabling SSH server..."
systemctl enable ssh
systemctl start ssh

header "Hostname Setup"
CURRENT_HOSTNAME=$(hostname)
echo -e "Current hostname: ${YELLOW}$CURRENT_HOSTNAME${NC}"
read -rp "Enter a hostname for this spoke (e.g. rocky, gouda, camembert): " SPOKE_NAME
SPOKE_NAME="${SPOKE_NAME:-$CURRENT_HOSTNAME}"
hostnamectl set-hostname "$SPOKE_NAME"
echo "$SPOKE_NAME" > /etc/hostname
info "Hostname set to: $SPOKE_NAME"

header "SSH Key Setup"
echo "How would you like to handle the SSH key for the tunnel to $HUB_HOST?"
echo "  1) Generate a new key automatically"
echo "  2) Use an existing key (paste the private key)"
echo ""
read -rp "Choose [1/2]: " KEY_CHOICE

case "$KEY_CHOICE" in
1)
    KEY_NAME="oilykey"
    KEY_PATH="$SSH_DIR/$KEY_NAME"
    mkdir -p "$SSH_DIR"
    chown armbian:armbian "$SSH_DIR"
    chmod 700 "$SSH_DIR"

    if [ -f "$KEY_PATH" ]; then
        warn "Key $KEY_PATH already exists, using it."
    else
        info "Generating new ED25519 key..."
        sudo -u armbian ssh-keygen -t ed25519 -f "$KEY_PATH" -N ""
        chown armbian:armbian "$KEY_PATH" "$KEY_PATH.pub"
        chmod 600 "$KEY_PATH"
    fi

    echo ""
    echo -e "${YELLOW}══════════════════════════════════════════${NC}"
    echo -e "${YELLOW}  Send this public key to finn (oily.dad owner)${NC}"
    echo -e "${YELLOW}  and ask him to add it to armbian@oily.dad authorized_keys:${NC}"
    echo -e "${YELLOW}══════════════════════════════════════════${NC}"
    cat "$KEY_PATH.pub"
    echo -e "${YELLOW}══════════════════════════════════════════${NC}"
    echo ""
    read -rp "Press ENTER once finn has added the key to oily.dad..."
    ;;
2)
    read -rp "Enter a name for the key file (e.g. oilykey): " KEY_NAME
    KEY_NAME="${KEY_NAME:-oilykey}"
    KEY_PATH="$SSH_DIR/$KEY_NAME"
    mkdir -p "$SSH_DIR"
    chown armbian:armbian "$SSH_DIR"
    chmod 700 "$SSH_DIR"

    echo "Paste the private key content below, then press ENTER and CTRL+D:"
    KEY_CONTENT=$(cat | tr -d '\r')
    printf '%s\n' "$KEY_CONTENT" > "$KEY_PATH"
    chown armbian:armbian "$KEY_PATH"
    chmod 600 "$KEY_PATH"
    info "Key saved to $KEY_PATH"
    ;;
*)
    die "Invalid choice"
    ;;
esac

info "Scanning hub host key..."
sudo -u armbian touch "$SSH_DIR/known_hosts"
chown armbian:armbian "$SSH_DIR/known_hosts"
chmod 600 "$SSH_DIR/known_hosts"
sudo -u armbian ssh-keyscan -H "$HUB_HOST" >> "$SSH_DIR/known_hosts" 2>/dev/null

header "Testing SSH Connection"
info "Testing connection to $HUB_HOST..."
retry_or_abort \
    "sudo -u armbian ssh -i \"$KEY_PATH\" -o BatchMode=yes -o ConnectTimeout=10 \"$HUB_USER@$HUB_HOST\" exit" \
    "SSH connection to $HUB_HOST failed. Check that finn added your public key."

header "Finding Available Tunnel Port"
info "Scanning for a free port on $HUB_HOST starting from $START_PORT..."
TUNNEL_PORT=""
for PORT in $(seq "$START_PORT" $((START_PORT + 20))); do
    RESULT=$(sudo -u armbian ssh -i "$KEY_PATH" "$HUB_USER@$HUB_HOST" "ss -tlnp | grep :$PORT" 2>/dev/null || true)
    if [ -z "$RESULT" ]; then
        TUNNEL_PORT="$PORT"
        info "Port $TUNNEL_PORT is available."
        break
    else
        warn "Port $PORT is in use, trying next..."
    fi
done

[ -n "$TUNNEL_PORT" ] || die "Could not find a free port between $START_PORT and $((START_PORT + 20)). Ask finn to free up a port."

header "Configuring compose.yaml"
info "Setting port to $TUNNEL_PORT and key to $KEY_NAME..."

sed -i "s|-R [0-9]*:localhost:22|-R ${TUNNEL_PORT}:localhost:22|g" "$COMPOSE"
sed -i "s|-i /home/armbian/.ssh/[^ ]*|-i /home/armbian/.ssh/${KEY_NAME}|g" "$COMPOSE"
sed -i "s|/home/armbian/.ssh/oilykey[^:]*:/home/armbian/.ssh/oilykey[^:]*|${SSH_DIR}/${KEY_NAME}:${SSH_DIR}/${KEY_NAME}|g" "$COMPOSE"
sed -i "s|container_name: spoke-autossh|container_name: ${SPOKE_NAME}-autossh|g" "$COMPOSE"
sed -i "s|container_name: spoke-syncthing|container_name: ${SPOKE_NAME}-syncthing|g" "$COMPOSE"
sed -i "s|hostname: spoke-syncthing|hostname: ${SPOKE_NAME}-syncthing|g" "$COMPOSE"
sed -i '/^version:/d' "$COMPOSE"

SYNCTHING_MOUNT="$ARMBIAN_HOME/st"
mkdir -p "$SYNCTHING_MOUNT"
chown armbian:armbian "$SYNCTHING_MOUNT"

header "Building Docker Image"
cd "$SPOKE_DIR"
docker build \
    --build-arg UID="$(id -u armbian)" \
    --build-arg GID="$(id -g armbian)" \
    -t spoke-autossh .

header "Starting Containers"
docker compose up -d
info "Waiting for tunnel to establish..."
sleep 6

LOGS=$(docker logs "${SPOKE_NAME}-autossh" 2>&1 || docker logs spoke-autossh 2>&1 || true)
if echo "$LOGS" | grep -q "remote port forwarding failed"; then
    warn "Tunnel failed — port $TUNNEL_PORT may have been taken between check and connect."
    warn "Try running: docker compose down && docker compose up -d"
    warn "Or re-run this script."
else
    info "Tunnel is up on port $TUNNEL_PORT."
fi

header "Setup Complete"
echo -e "  Spoke name:    ${GREEN}$SPOKE_NAME${NC}"
echo -e "  Tunnel port:   ${GREEN}$TUNNEL_PORT${NC} on $HUB_HOST"
echo -e "  SSH key:       ${GREEN}$KEY_PATH${NC}"
echo ""
echo -e "${YELLOW}Finn needs to do the following on oily.dad:${NC}"
echo ""
echo "  1. Generate a hub->spoke key:"
echo "     ssh-keygen -t ed25519 -f ~/.ssh/armbian-${SPOKE_NAME}-$(date +%Y%m)"
echo ""
echo "  2. Copy it to this spoke through the tunnel:"
echo "     ssh-copy-id -i ~/.ssh/armbian-${SPOKE_NAME}-$(date +%Y%m).pub -p $TUNNEL_PORT armbian@localhost"
echo ""
echo "  3. Add an rclone remote in ~/.config/rclone/rclone.conf:"
echo "     [${SPOKE_NAME}-remote]"
echo "     type = sftp"
echo "     host = localhost"
echo "     port = $TUNNEL_PORT"
echo "     key_file = /home/armbian/.ssh/armbian-${SPOKE_NAME}-$(date +%Y%m)"
echo "     shell_type = unix"
echo "     md5sum_command = md5sum"
echo "     sha1sum_command = sha1sum"
echo ""