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

# ─────────────────────────────────────────────────────────────────────────────
# One-shot server bootstrap: fresh Arch VPS → hardened box → native OMATERM.
#
# Run as root over SSH **with a TTY** (the install ends in an interactive
# omaterm session, so a plain `ssh host cmd` won't do):
#
#   ssh -t root@YOUR_SERVER 'bash <(curl -fsSL https://raw.githubusercontent.com/28allday/omaserver/main/install-server.sh)'
#
# Phase A (as root):   hardening — deploy user, key-only SSH, ufw, fail2ban,
#                      weekly updates, swap, journald cap, cache hygiene
# Phase B (as user):   hands off to install-native.sh from the frozen fork
#                      (28allday/omaterm @ native-frozen — the branch omalocal
#                      also pins; fetched at runtime, never modified) which
#                      installs packages, omadots, AI agents, and runs the
#                      interactive first-time setup
#
# Safe to re-run — completed steps skip themselves. After the first run root
# SSH login is disabled, so re-run as the deploy user with sudo.
# ─────────────────────────────────────────────────────────────────────────────

# =============================================================================
# DEFAULTS
# =============================================================================

TIMEZONE="${TIMEZONE:-Europe/London}"
LOCALE="${LOCALE:-en_GB.UTF-8}"

# Where install-native.sh is fetched from at runtime — the frozen branch that
# omalocal also pins. Read-only consumer: nothing is ever pushed there.
OMATERM_REPO_RAW="https://raw.githubusercontent.com/28allday/omaterm/native-frozen"

# This script's own home, for re-run instructions
SELF_URL="https://raw.githubusercontent.com/28allday/omaserver/main/install-server.sh"

# =============================================================================
# COLOURS
# =============================================================================

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

log()  { echo -e "${GREEN}[✓]${NC} $1"; }
warn() { echo -e "${YELLOW}[!]${NC} $1"; }
err()  { echo -e "${RED}[✗]${NC} $1"; }
step() { echo -e "\n${CYAN}━━━ $1 ━━━${NC}\n"; }

# =============================================================================
# PRE-FLIGHT CHECKS
# =============================================================================

if [ "$EUID" -ne 0 ]; then
    err "This script must be run as root."
    err "Usage: ssh -t root@YOUR_SERVER 'bash <(curl -fsSL $SELF_URL)'"
    exit 1
fi

if [ ! -f /etc/arch-release ]; then
    err "This script requires Arch Linux."
    exit 1
fi

if [ ! -t 0 ]; then
    err "No TTY on stdin — the omaterm installer needs an interactive terminal."
    err "Connect with:  ssh -t root@YOUR_SERVER"
    err "Then run:      bash <(curl -fsSL $SELF_URL)"
    exit 1
fi

# =============================================================================
# INTERACTIVE SETUP
# =============================================================================

echo -e "${CYAN}"
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║           OMATERM Server Bootstrap                           ║"
echo "║           Arch Linux · Hardening · Native OMATERM            ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo -e "${NC}"

# Public IP — used in re-run instructions and the SSH config snippet
SERVER_IP=$(curl -4 -s --max-time 5 ifconfig.me 2>/dev/null || \
            curl -4 -s --max-time 5 icanhazip.com 2>/dev/null || \
            echo "YOUR_SERVER_IP")

# Answers survive the kernel-upgrade reboot — a re-run finds them here and
# doesn't ask again
STATE_FILE="/root/.install-server.conf"
SSH_ALIAS=""
DEPLOY_USER=""
GIT_NAME=""
GIT_EMAIL=""
if [ -f "$STATE_FILE" ]; then
    # shellcheck source=/dev/null
    . "$STATE_FILE"
fi

if [ -n "$SSH_ALIAS" ] && [ -n "$DEPLOY_USER" ]; then
    log "Using saved settings from the previous run"
else
    # Machine name — used as the hostname, the tailnet name, AND the alias in
    # your local SSH config (ssh <name>). Asked ONCE, used everywhere.
    echo ""
    echo "  This name becomes the hostname, the Tailscale machine name,"
    echo "  and your local SSH alias (ssh <name>)."
    read -rp "Machine name for this server [vps]: " input_alias
    SSH_ALIAS="${input_alias:-vps}"

    # Username
    echo ""
    read -rp "Choose a username for the deploy user [deploy]: " input_user
    DEPLOY_USER="${input_user:-deploy}"

    # Git identity — pre-seeds the deploy user's ~/.gitconfig so omaterm's
    # first-run setup skips its git prompts. Leave blank to be asked there
    # instead.
    echo ""
    read -rp "Git name (blank to skip): " GIT_NAME
    if [ -n "$GIT_NAME" ]; then
        read -rp "Git email: " GIT_EMAIL
    fi
fi

# Validate username
if ! [[ "$DEPLOY_USER" =~ ^[a-z_][a-z0-9_-]*$ ]]; then
    err "Invalid username. Use lowercase letters, numbers, hyphens, underscores."
    rm -f "$STATE_FILE"
    exit 1
fi

# Confirm
echo ""
echo -e "${CYAN}━━━ Confirm settings ━━━${NC}"
echo ""
echo "  Machine name:   $SSH_ALIAS  (hostname + tailnet + ssh alias)"
echo "  Deploy user:    $DEPLOY_USER"
echo "  Git identity:   ${GIT_NAME:-—} ${GIT_EMAIL:+<$GIT_EMAIL>}"
echo "  Timezone:       $TIMEZONE"
echo "  Omaterm:        native (28allday/omaterm @ native-frozen)"
echo ""
read -p "Continue? (Y/n) " -n 1 -r
echo ""
if [[ $REPLY =~ ^[Nn]$ ]]; then
    rm -f "$STATE_FILE"
    echo "Aborted — saved settings cleared. Re-run to start over."
    exit 0
fi

# Persist for re-runs (cleared on abort above)
printf 'SSH_ALIAS=%q\nDEPLOY_USER=%q\nGIT_NAME=%q\nGIT_EMAIL=%q\n' \
    "$SSH_ALIAS" "$DEPLOY_USER" "$GIT_NAME" "$GIT_EMAIL" > "$STATE_FILE"
chmod 600 "$STATE_FILE"

# =============================================================================
# 1. SYSTEM UPDATES & ESSENTIALS
# =============================================================================

step "1/9 · System updates & essential packages"

pacman -Syu --noconfirm

# If -Syu upgraded the kernel, the running kernel's modules are gone and
# things like tun (tailscale) and netfilter (docker/ufw) can't load —
# reboot and re-run (script is safe to re-run; completed steps skip
# themselves). This runs BEFORE SSH hardening, so root login still works
# for the re-run.
if [ ! -d "/usr/lib/modules/$(uname -r)" ]; then
    warn "Kernel was upgraded — running kernel $(uname -r) has no modules on disk."
    warn "The box must reboot before setup can continue."
    echo ""
    echo "  After the reboot, reconnect and re-run — it picks up where it left off:"
    echo ""
    if [ -f "$0" ] && [[ "$0" != /dev/fd/* ]]; then
        echo -e "    ${CYAN}ssh -t root@$SERVER_IP bash $(basename "$0")${NC}"
    else
        echo -e "    ${CYAN}ssh -t root@$SERVER_IP 'bash <(curl -fsSL $SELF_URL)'${NC}"
    fi
    echo ""
    read -p "Reboot now? (Y/n) " -n 1 -r
    echo ""
    if [[ ! $REPLY =~ ^[Nn]$ ]]; then
        systemctl reboot
    fi
    exit 1
fi

# Hardening-phase packages only — install-native.sh brings everything else
# (docker, tailscale, git, tmux, agents) via arch.packages.
# kitty-terminfo: without it, SSHing in from kitty gives
# "'xterm-kitty': unknown terminal type" and interactive tools die
pacman -S --noconfirm --needed \
    ufw \
    fail2ban \
    curl \
    kitty-terminfo \
    pacman-contrib

log "System updated, hardening packages installed"

# =============================================================================
# 2. TIMEZONE & LOCALE
# =============================================================================

step "2/9 · Hostname, timezone & locale"

# Machine name becomes the hostname — omaterm-setup's tailnet prompt
# defaults to it, so the name is only ever typed once
hostnamectl set-hostname "$SSH_ALIAS"
log "Hostname set to $SSH_ALIAS"

timedatectl set-timezone "$TIMEZONE"
log "Timezone set to $TIMEZONE"

# Generate and set locale (minimal Arch cloud images often ship C.UTF-8 only)
if ! locale -a 2>/dev/null | grep -qi "^${LOCALE/UTF-8/utf8}$"; then
    sed -i "s/^#${LOCALE}/${LOCALE}/" /etc/locale.gen
    locale-gen
fi
echo "LANG=$LOCALE" > /etc/locale.conf
log "Locale set to $LOCALE"

# =============================================================================
# 3. CREATE DEPLOY USER
# =============================================================================

step "3/9 · Deploy user"

if id "$DEPLOY_USER" &>/dev/null; then
    warn "User '$DEPLOY_USER' already exists — skipping creation"
else
    useradd -m -s /bin/bash "$DEPLOY_USER"
    usermod -aG wheel "$DEPLOY_USER"
    log "User '$DEPLOY_USER' created and added to wheel group"
fi

# Passwordless sudo — install-native.sh and omaterm-setup sudo freely,
# and the user has no password to type
echo "$DEPLOY_USER ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/"$DEPLOY_USER"
chmod 440 /etc/sudoers.d/"$DEPLOY_USER"
log "Passwordless sudo enabled for $DEPLOY_USER"

# Pre-seed git identity — omaterm-setup skips its git prompts when
# user.name and user.email are already configured
if [ -n "$GIT_NAME" ] && [ ! -f /home/"$DEPLOY_USER"/.gitconfig ]; then
    cat > /home/"$DEPLOY_USER"/.gitconfig << GITEOF
[user]
	name = $GIT_NAME
	email = $GIT_EMAIL
GITEOF
    chown "$DEPLOY_USER":"$DEPLOY_USER" /home/"$DEPLOY_USER"/.gitconfig
    log "Git identity pre-seeded for $DEPLOY_USER"
fi

# Copy root's authorized_keys to deploy user if they exist
if [ -f /root/.ssh/authorized_keys ]; then
    mkdir -p /home/"$DEPLOY_USER"/.ssh
    cp /root/.ssh/authorized_keys /home/"$DEPLOY_USER"/.ssh/authorized_keys
    chown -R "$DEPLOY_USER":"$DEPLOY_USER" /home/"$DEPLOY_USER"/.ssh
    chmod 700 /home/"$DEPLOY_USER"/.ssh
    chmod 600 /home/"$DEPLOY_USER"/.ssh/authorized_keys
    log "SSH keys copied from root to $DEPLOY_USER"
else
    warn "No SSH keys found for root — you'll need to add keys manually:"
    warn "  ssh-copy-id $DEPLOY_USER@YOUR_SERVER_IP"
fi

# =============================================================================
# 4. HARDEN SSH
# =============================================================================

step "4/9 · SSH hardening"

# Backup original config
cp /etc/ssh/sshd_config "/etc/ssh/sshd_config.backup.$(date +%Y%m%d)"

# Neutralise conflicting directives in the main config so our drop-in wins
sed -i 's/^\s*PermitRootLogin\b/#&/'          /etc/ssh/sshd_config
sed -i 's/^\s*PasswordAuthentication\b/#&/'    /etc/ssh/sshd_config
sed -i 's/^\s*PubkeyAuthentication\b/#&/'      /etc/ssh/sshd_config
sed -i 's/^\s*X11Forwarding\b/#&/'             /etc/ssh/sshd_config
sed -i 's/^\s*PermitEmptyPasswords\b/#&/'      /etc/ssh/sshd_config
log "Commented out conflicting directives in sshd_config"

# Remove cloud-init SSH overrides that can conflict (some VPS images use them)
rm -f /etc/ssh/sshd_config.d/50-cloud-init.conf

# Ensure the drop-in directory is included (Arch's stock sshd_config has this,
# but some provider images strip it)
mkdir -p /etc/ssh/sshd_config.d
if ! grep -q '^Include /etc/ssh/sshd_config.d/\*\.conf' /etc/ssh/sshd_config; then
    sed -i '1i Include /etc/ssh/sshd_config.d/*.conf' /etc/ssh/sshd_config
    warn "Added missing Include for sshd_config.d drop-ins"
fi

cat > /etc/ssh/sshd_config.d/hardened.conf << EOF
# Custom SSH hardening
PermitRootLogin no
PasswordAuthentication no
PubkeyAuthentication yes
MaxAuthTries 3
MaxSessions 5
AllowUsers $DEPLOY_USER
ClientAliveInterval 300
ClientAliveCountMax 2
X11Forwarding no
PermitEmptyPasswords no
EOF

log "SSH hardened — root disabled, key-only auth"

# Test config before restarting (service is 'sshd' on Arch, not 'ssh')
if sshd -t; then
    systemctl restart sshd
    log "SSH restarted successfully"
else
    err "SSH config test failed — reverting"
    rm /etc/ssh/sshd_config.d/hardened.conf
    exit 1
fi

# =============================================================================
# 5. FIREWALL (UFW)
# =============================================================================

step "5/9 · Firewall (UFW)"

ufw default deny incoming
ufw default allow outgoing
ufw allow 22/tcp comment 'SSH'

# Native omaterm runs tailscaled on the host (kernel networking, tun
# device) — without this, incoming tailnet traffic is silently dropped
ufw allow in on tailscale0 comment 'Tailscale'

# Enable without prompt, and make it survive reboots
echo "y" | ufw enable
systemctl enable ufw

log "UFW enabled — 22/tcp + tailscale0 only"
ufw status verbose

# =============================================================================
# 6. FAIL2BAN
# =============================================================================

step "6/9 · Fail2ban"

# Arch logs auth to the systemd journal — there is no /var/log/auth.log,
# so use the systemd backend instead of a logpath
cat > /etc/fail2ban/jail.local << EOF
[DEFAULT]
bantime = 3600
findtime = 600
maxretry = 3
banaction = ufw
backend = systemd

[sshd]
enabled = true
port = 22
filter = sshd
maxretry = 3
bantime = 3600
EOF

systemctl enable fail2ban
systemctl restart fail2ban
log "Fail2ban configured and running (systemd journal backend)"

# =============================================================================
# 7. AUTOMATIC UPDATES (weekly pacman -Syu timer)
# =============================================================================

step "7/9 · Automatic updates"

# Arch has no unattended-upgrades equivalent — use a weekly full update timer.
# NOTE: rolling release means an update can very occasionally need manual
# intervention (check https://archlinux.org/news/ if something breaks).
cat > /etc/systemd/system/pacman-update.service << EOF
[Unit]
Description=Weekly pacman system update
After=network-online.target
Wants=network-online.target

[Service]
Type=oneshot
ExecStart=/usr/bin/pacman -Syu --noconfirm
EOF

cat > /etc/systemd/system/pacman-update.timer << EOF
[Unit]
Description=Weekly pacman system update

[Timer]
OnCalendar=weekly
RandomizedDelaySec=1h
Persistent=true

[Install]
WantedBy=timers.target
EOF

systemctl daemon-reload
systemctl enable --now pacman-update.timer
log "Weekly auto-update timer enabled (pacman-update.timer)"
warn "Rolling release: check 'journalctl -u pacman-update' if anything misbehaves after an update"

# =============================================================================
# 8. HOUSEKEEPING — swap, services, journal, pacman cache
# =============================================================================

step "8/9 · Housekeeping"

# --- Swap (2GB) ---
if swapon --show --noheadings 2>/dev/null | grep -q .; then
    warn "Swap already active — skipping:"
    swapon --show
elif [ -f /swapfile ]; then
    warn "Swap file already exists — skipping"
else
    # fallocate swapfiles fail on btrfs (needs NOCOW) — use the proper tool
    if [ "$(findmnt -n -o FSTYPE /)" = "btrfs" ]; then
        btrfs filesystem mkswapfile --size 2g /swapfile
    else
        dd if=/dev/zero of=/swapfile bs=1M count=2048 status=none
        chmod 600 /swapfile
        mkswap /swapfile
    fi
    swapon /swapfile
    grep -q '/swapfile' /etc/fstab || echo '/swapfile none swap sw 0 0' >> /etc/fstab
    log "2GB swap created and enabled"
fi

# --- Disable unnecessary services ---
for svc in ModemManager udisks2 multipathd; do
    if systemctl is-enabled "$svc" &>/dev/null; then
        systemctl stop "$svc"
        systemctl disable "$svc"
        log "Disabled $svc"
    fi
done

# --- Cap journal logs (50MB) ---
mkdir -p /etc/systemd/journald.conf.d
cat > /etc/systemd/journald.conf.d/size.conf << EOF
[Journal]
SystemMaxUse=50M
EOF
systemctl restart systemd-journald
log "Journal capped at 50MB"

# --- Pacman cache hygiene (keep last versions, prune weekly) ---
systemctl enable --now paccache.timer
log "paccache.timer enabled — old package versions pruned weekly"

# =============================================================================
# 9. OMATERM HANDOFF
# =============================================================================

step "9/9 · Omaterm handoff"

# Pre-seed the tailnet machine name — omaterm-setup's hostname prompt
# defaults to this file's contents, so Enter accepts the machine name
echo "$SSH_ALIAS" > /etc/omaterm-tailscale-hostname

# Re-runnable installer wrapper in the deploy user's home. Download-then-run
# (not curl | bash) so stdin stays on the real TTY — the installer's
# interactive first-run setup (gum) and final shell handoff need it.
cat > /home/"$DEPLOY_USER"/install-omaterm << OMAEOF
#!/bin/bash
# Fetch and run the omaterm native installer (28allday fork, native-frozen)
# with a real TTY on stdin. Re-run this if the install is interrupted
# (e.g. after a kernel-upgrade reboot).
set -uo pipefail
tmp=\$(mktemp)
trap 'rm -f "\$tmp"' EXIT
curl -fsSL $OMATERM_REPO_RAW/install-native.sh -o "\$tmp" || exit 1

OMATERM_REF=native-frozen bash "\$tmp"

# On success install-native.sh execs into a login shell and never returns —
# reaching this point means the chain bailed early (a first-run setup
# section erroring out, a stray Ctrl+C). Don't strand the user by dropping
# the SSH session: resume setup if it didn't finish, then land in a shell.
echo ""
echo "[!] The installer exited early — recovering."
if [ ! -f "\$HOME/.omaterm-setup-done" ] && [ -x "\$HOME/.local/bin/omaterm-setup" ]; then
    echo "[!] Resuming first-run setup (already-completed sections are skipped)..."
    "\$HOME/.local/bin/omaterm-setup" || true
fi
exec bash -l
OMAEOF

chown "$DEPLOY_USER":"$DEPLOY_USER" /home/"$DEPLOY_USER"/install-omaterm
chmod +x /home/"$DEPLOY_USER"/install-omaterm
log "Installer wrapper created at ~$DEPLOY_USER/install-omaterm"

# =============================================================================
# LOCAL SSH CONFIG SNIPPET
# =============================================================================

cat > /home/"$DEPLOY_USER"/ssh-config-snippet.txt << SSHEOF
Host $SSH_ALIAS
    HostName $SERVER_IP
    User $DEPLOY_USER
    IdentityFile ~/.ssh/id_ed25519
SSHEOF

chown "$DEPLOY_USER":"$DEPLOY_USER" /home/"$DEPLOY_USER"/ssh-config-snippet.txt

# Hardening phase complete — saved answers no longer needed
rm -f "$STATE_FILE"

# =============================================================================
# SUMMARY
# =============================================================================

echo ""
echo -e "${CYAN}"
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║                  HARDENING COMPLETE                          ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo -e "${NC}"
echo ""
echo -e "  ${GREEN}SSH${NC}"
echo "  ├─ User:          $DEPLOY_USER"
echo "  ├─ Root login:    disabled"
echo "  └─ Auth:          key-only"
echo ""
echo -e "  ${GREEN}Firewall${NC}"
echo "  ├─ 22/tcp         SSH"
echo "  └─ tailscale0     allowed (for omaterm's tailscale up --ssh)"
echo ""
echo -e "  ${GREEN}Services${NC}"
echo "  ├─ Fail2ban:      active (systemd backend)"
echo "  └─ Auto-updates:  weekly pacman-update.timer"
echo ""
echo "  1. On your local machine, add the SSH config:"
echo ""
echo -e "  ${CYAN}────── COPY BELOW THIS LINE ──────${NC}"
echo ""
echo "cat >> ~/.ssh/config << 'EOF'"
echo ""
echo "Host $SSH_ALIAS"
echo "    HostName $SERVER_IP"
echo "    User $DEPLOY_USER"
echo "    IdentityFile ~/.ssh/id_ed25519"
echo "EOF"
echo ""
echo -e "  ${CYAN}────── COPY ABOVE THIS LINE ──────${NC}"
echo ""
echo "  2. The omaterm install runs next, as $DEPLOY_USER. If it's"
echo "     interrupted (e.g. a kernel-upgrade reboot), resume with:"
echo ""
echo -e "       ${CYAN}ssh -t $SSH_ALIAS ./install-omaterm${NC}"
echo ""
echo -e "  ${YELLOW}Heads-up — after the packages install, first-run setup asks about:${NC}"
echo ""
if [ -n "$GIT_NAME" ]; then
    echo "     · Git identity            skipped — already configured"
else
    echo "     · Git identity            your name + email"
fi
echo "     · GitHub                  optional, wants a classic token"
echo "     · Tailscale               recommended — browser login link;"
echo "                               hostname pre-filled ($SSH_ALIAS), Enter accepts"
echo -e "     · SSH public key         ${YELLOW}answer No${NC} — this box is already key-only"
echo "     · 1Password               optional"
echo ""
echo -e "  ${YELLOW}A reboot is recommended once you're done, to apply kernel updates.${NC}"
echo ""

read -p "Launch the omaterm installer now as $DEPLOY_USER? (Y/n) " -n 1 -r
echo ""
if [[ ! $REPLY =~ ^[Nn]$ ]]; then
    # su keeps the controlling TTY, so gum prompts and the final
    # `exec bash -l` in install-native.sh all work
    exec su - "$DEPLOY_USER" -c ./install-omaterm
else
    echo ""
    echo "  When you're ready:"
    echo ""
    echo -e "       ${CYAN}ssh -t $SSH_ALIAS ./install-omaterm${NC}"
    echo ""
fi