#!/usr/bin/env bash
# Run this script on the LOCAL machine (where Ollama runs).
# It registers your SSH public key on the server and installs
# the autossh systemd service for a persistent tunnel.
#
# Usage: ./setup.sh <server-ip-or-hostname>
#
# Prerequisites (local machine):
#   apt/brew: openssh-client autossh
#   Ollama running on localhost:11434

set -euo pipefail

SERVER="${1:?Usage: $0 <server-ip-or-hostname>}"
TUNNEL_USER="ollama-tunnel"
REMOTE_PORT=11434
LOCAL_PORT=11434
KEY_FILE="${HOME}/.ssh/id_ed25519"

# Generate key if it doesn't exist
if [[ ! -f "${KEY_FILE}" ]]; then
    echo "[+] Generating SSH key ${KEY_FILE} ..."
    ssh-keygen -t ed25519 -f "${KEY_FILE}" -N "" -C "ollama-tunnel@$(hostname)"
fi

# Copy public key to server
echo "[+] Copying public key to ${TUNNEL_USER}@${SERVER} ..."
echo "    You will be prompted for sudo on the server (or use the superseller account)."
PUBKEY=$(cat "${KEY_FILE}.pub")
ssh superseller@"${SERVER}" "sudo bash -c 'echo \"${PUBKEY}\" >> /home/${TUNNEL_USER}/.ssh/authorized_keys && sort -u /home/${TUNNEL_USER}/.ssh/authorized_keys -o /home/${TUNNEL_USER}/.ssh/authorized_keys'"

echo "[+] Testing tunnel connection ..."
ssh -o StrictHostKeyChecking=accept-new -o ConnectTimeout=5 \
    -N -i "${KEY_FILE}" \
    -R "172.18.0.1:${REMOTE_PORT}:localhost:${LOCAL_PORT}" \
    "${TUNNEL_USER}@${SERVER}" &
SSH_PID=$!
sleep 2
if kill -0 "${SSH_PID}" 2>/dev/null; then
    echo "[+] Tunnel works! Stopping test connection."
    kill "${SSH_PID}"
else
    echo "[!] Tunnel test failed. Check sshd config and firewall on the server."
    exit 1
fi

# Install systemd service
install_systemd_service() {
    local service_file="${HOME}/.config/systemd/user/ollama-tunnel.service"
    mkdir -p "$(dirname "${service_file}")"
    cat > "${service_file}" << EOF
[Unit]
Description=Ollama SSH reverse tunnel to SuperSeller3000 server
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
ExecStart=/usr/bin/autossh -M 0 \\
    -o "ServerAliveInterval=30" \\
    -o "ServerAliveCountMax=3" \\
    -o "ExitOnForwardFailure=yes" \\
    -o "StrictHostKeyChecking=accept-new" \\
    -N -i ${KEY_FILE} \\
    -R 172.18.0.1:${REMOTE_PORT}:localhost:${LOCAL_PORT} \\
    ${TUNNEL_USER}@${SERVER}
Restart=always
RestartSec=10

[Install]
WantedBy=default.target
EOF

    systemctl --user daemon-reload
    systemctl --user enable ollama-tunnel.service
    systemctl --user start ollama-tunnel.service
    echo "[+] systemd service installed and started."
    echo "    Status: systemctl --user status ollama-tunnel"
}

if command -v autossh &>/dev/null && command -v systemctl &>/dev/null; then
    echo "[+] Installing autossh systemd user service ..."
    install_systemd_service
else
    echo "[!] autossh or systemd not found. Manual tunnel command:"
    echo ""
    echo "    autossh -M 0 -o ServerAliveInterval=30 -N \\"
    echo "        -i ${KEY_FILE} \\"
    echo "        -R 172.18.0.1:${REMOTE_PORT}:localhost:${LOCAL_PORT} \\"
    echo "        ${TUNNEL_USER}@${SERVER}"
fi

echo ""
echo "Done. The server will see Ollama at http://172.18.0.1:${REMOTE_PORT}"