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

HUB_HOST=""
HUB_USER=""
SPOKE_USER=""
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SPOKE_DIR="$SCRIPT_DIR"
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'

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

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
}

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
}

check_permissions() {
    local file="$1"
    local label="$2"
    if [ ! -f "$file" ]; then
        warn "Permission check: $label not found at $file"
        return
    fi
    local perms
    perms=$(stat -c "%a" "$file" 2>/dev/null || stat -f "%OLp" "$file" 2>/dev/null)
    if [ -z "$perms" ]; then
        warn "Could not read permissions for $label ($file)"
        return
    fi
    local world="${perms: -1}"
    local group="${perms: -2:1}"
    if [ "$world" != "0" ] || [ "$group" != "0" ]; then
        warn "UNSAFE PERMISSIONS on $label ($file): $perms — should be 600 or 400"
        warn "Fixing permissions automatically..."
        chmod 600 "$file"
        info "Permissions fixed: $file is now 600"
    else
        info "Permissions OK: $label ($file) = $perms"
    fi
}

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

check_deps ip ssh ssh-keygen ssh-keyscan systemctl hostnamectl

read -rp "Hub hostname [oily.dad]: " HUB_HOST
HUB_HOST="${HUB_HOST:-oily.dad}"
read -rp "Hub SSH user [armbian]: " HUB_USER
HUB_USER="${HUB_USER:-armbian}"
read -rp "Spoke local user [armbian]: " SPOKE_USER
SPOKE_USER="${SPOKE_USER:-armbian}"
ARMBIAN_HOME="/home/$SPOKE_USER"
SSH_DIR="$ARMBIAN_HOME/.ssh"

header "TinyBoard Spoke Setup"

header "Detecting Package Manager"
if command -v apt-get >/dev/null 2>&1; then
    PKG_MANAGER="apt"
    PKG_INSTALL="apt-get install -y -q"
    OPENSSH_PKG="openssh-server"
    AUTOSSH_PKG="autossh"
    info "Detected: apt (Debian/Ubuntu)"
    apt-get update -q
elif command -v dnf >/dev/null 2>&1; then
    PKG_MANAGER="dnf"
    PKG_INSTALL="dnf install -y -q"
    OPENSSH_PKG="openssh-server"
    AUTOSSH_PKG="autossh"
    info "Detected: dnf (Fedora/RHEL/Alma/Rocky)"
    dnf check-update -q || true
elif command -v yum >/dev/null 2>&1; then
    PKG_MANAGER="yum"
    PKG_INSTALL="yum install -y -q"
    OPENSSH_PKG="openssh-server"
    AUTOSSH_PKG="autossh"
    info "Detected: yum (older RHEL/CentOS)"
    yum check-update -q || true
elif command -v pacman >/dev/null 2>&1; then
    PKG_MANAGER="pacman"
    PKG_INSTALL="pacman -S --noconfirm --quiet"
    OPENSSH_PKG="openssh"
    AUTOSSH_PKG="autossh"
    info "Detected: pacman (Arch)"
    pacman -Sy --quiet
else
    die "No supported package manager found (apt, dnf, yum, pacman)"
fi

header "Installing Packages"
if ! command -v curl >/dev/null 2>&1; then
    $PKG_INSTALL curl
fi

$PKG_INSTALL vim "$AUTOSSH_PKG" "$OPENSSH_PKG" git

info "Installing Docker..."
if ! command -v docker >/dev/null 2>&1; then
    if [ "$PKG_MANAGER" = "apt" ]; then
        $PKG_INSTALL docker.io docker-compose-plugin
    else
        curl -fsSL https://get.docker.com | bash
    fi
else
    warn "Docker already installed, skipping."
fi

if ! docker compose version >/dev/null 2>&1; then
    if [ "$PKG_MANAGER" = "apt" ]; then
        $PKG_INSTALL docker-compose-plugin
    else
        warn "docker compose not available — Docker install script should have included it."
    fi
fi

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

info "Enabling SSH server..."
if systemctl enable ssh 2>/dev/null; then
    systemctl start ssh
elif systemctl enable sshd 2>/dev/null; then
    systemctl start sshd
else
    warn "Could not enable SSH service — please start it manually."
fi

SSHD_CONF="/etc/ssh/sshd_config"

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)
    read -rp "Key name [hubkey]: " KEY_NAME
    KEY_NAME="${KEY_NAME:-hubkey}"
    KEY_PATH="$SSH_DIR/$KEY_NAME"
    mkdir -p "$SSH_DIR"
    chown "$SPOKE_USER":"$SPOKE_USER" "$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 "$SPOKE_USER" ssh-keygen -t ed25519 -f "$KEY_PATH" -N ""
        chown "$SPOKE_USER":"$SPOKE_USER" "$KEY_PATH" "$KEY_PATH.pub"
        chmod 600 "$KEY_PATH"
    fi

    echo ""
    echo -e "${YELLOW}══════════════════════════════════════════${NC}"
    echo -e "${YELLOW}  Send this public key to the hub owner${NC}"
    echo -e "${YELLOW}  and ask them to add it to ${HUB_USER}@${HUB_HOST} authorized_keys:${NC}"
    echo -e "${YELLOW}══════════════════════════════════════════${NC}"
    cat "$KEY_PATH.pub"
    echo -e "${YELLOW}══════════════════════════════════════════${NC}"
    echo ""
    read -rp "Press ENTER once the key has been added to ${HUB_HOST}..."
    ;;
2)
    read -rp "Enter a name for the key file [hubkey]: " KEY_NAME
    KEY_NAME="${KEY_NAME:-hubkey}"
    KEY_PATH="$SSH_DIR/$KEY_NAME"
    mkdir -p "$SSH_DIR"
    chown "$SPOKE_USER":"$SPOKE_USER" "$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 "$SPOKE_USER":"$SPOKE_USER" "$KEY_PATH"
    chmod 600 "$KEY_PATH"
    info "Key saved to $KEY_PATH"
    ;;
*)
    die "Invalid choice"
    ;;
esac

header "Password Authentication"
read -rp "Disable password auth for $SPOKE_USER and use keys only? [Y/n]: " DISABLE_PASS
DISABLE_PASS="${DISABLE_PASS:-y}"
if [[ "${DISABLE_PASS,,}" == "y" ]]; then
    if [ ! -f "$KEY_PATH" ]; then
        warn "No key found at $KEY_PATH — skipping password auth disable to avoid lockout."
    else
        if grep -q "^PasswordAuthentication" "$SSHD_CONF"; then
            sed -i "s/^PasswordAuthentication.*/PasswordAuthentication no/" "$SSHD_CONF"
        else
            echo "PasswordAuthentication no" >> "$SSHD_CONF"
        fi
        if grep -q "^PubkeyAuthentication" "$SSHD_CONF"; then
            sed -i "s/^PubkeyAuthentication.*/PubkeyAuthentication yes/" "$SSHD_CONF"
        else
            echo "PubkeyAuthentication yes" >> "$SSHD_CONF"
        fi
        info "Password authentication disabled for $SPOKE_USER."
        echo ""
        warn "Restarting SSH will apply the new settings."
        warn "If you are connected via SSH, your session may drop."
        warn "Make sure you can reconnect using your key before continuing."
        read -rp "Press ENTER to restart SSH or CTRL+C to abort..."
        if systemctl restart ssh 2>/dev/null; then
            info "SSH restarted."
        elif systemctl restart sshd 2>/dev/null; then
            info "SSH restarted."
        else
            warn "Could not restart SSH — please restart it manually."
        fi
    fi
else
    info "Password authentication left enabled."
fi

info "Checking SSH key permissions..."
check_permissions "$KEY_PATH" "spoke SSH private key"
if [ -f "$KEY_PATH.pub" ]; then
    check_permissions "$KEY_PATH.pub" "spoke SSH public key"
fi

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

header "Testing SSH Connection"
info "Testing connection to $HUB_HOST..."
retry_or_abort \
    "sudo -u \"$SPOKE_USER\" ssh -i \"$KEY_PATH\" -o BatchMode=yes -o ConnectTimeout=10 \"$HUB_USER@$HUB_HOST\" exit" \
    "SSH connection to $HUB_HOST failed. Check that the hub owner 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 "$SPOKE_USER" 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 the hub owner 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/[^ ]*/\.ssh/[^ ]*|-i ${SSH_DIR}/${KEY_NAME}|g" "$COMPOSE"
sed -i "s|/home/[^/]*/\.ssh/[^:]*:/home/[^/]*/\.ssh/[^:]*|${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 "$SPOKE_USER":"$SPOKE_USER" "$SYNCTHING_MOUNT"

header "Building Docker Image"
cd "$SPOKE_DIR"
docker build \
    --build-arg UID="$(id -u "$SPOKE_USER")" \
    --build-arg GID="$(id -g "$SPOKE_USER")" \
    -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}The hub owner needs to do the following on ${HUB_HOST}:${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 ${HUB_USER}@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/$HUB_USER/.ssh/armbian-${SPOKE_NAME}-$(date +%Y%m)"
echo "     shell_type = unix"
echo "     md5sum_command = md5sum"
echo "     sha1sum_command = sha1sum"
echo ""