#!/bin/bash
# =============================================================================
# VPS Setup & Hardening Script for Arch Linux
# For use with ONCE (Basecamp) + Docker deployments
# Run as root on a fresh Arch Linux VPS
# =============================================================================

set -euo pipefail

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

GITHUB_PAT=""
TIMEZONE="Europe/London"
LOCALE="en_GB.UTF-8"

# Comma-separated list of repos to clone (leave empty to skip)
# Format: "user/repo:/opt/apps/dirname,user/repo2:/opt/apps/dirname2"
REPOS_TO_CLONE=""

# =============================================================================
# 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: sudo bash arch-setup.sh"
    exit 1
fi

if ! grep -qi '^ID=arch' /etc/os-release 2>/dev/null; then
    warn "This script is designed for Arch Linux. Proceed with caution."
fi

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

echo -e "${CYAN}"
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║           VPS Setup & Hardening Script                      ║"
echo "║           Arch Linux · Docker · ONCE · Git                  ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo -e "${NC}"

# SSH alias
echo ""
echo "  This name is used in your local SSH config so you can"
echo "  connect with: ssh <alias>  (e.g. ssh titan, ssh nebula)"
read -rp "SSH alias 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}"

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

# Git name
read -rp "Git display name [Gavin]: " input_git_name
GIT_NAME="${input_git_name:-Gavin}"

# Git email
read -rp "Git email: " input_git_email
GIT_EMAIL="${input_git_email}"

if [ -z "$GIT_EMAIL" ]; then
    warn "No Git email set — you can configure this later"
    GIT_EMAIL="not-set@example.com"
fi

# GitHub PAT
echo ""
echo "  A GitHub Personal Access Token lets the server pull from"
echo "  your private repos and container registry (GHCR)."
echo "  You can skip this and add one later."
read -rp "GitHub username (leave empty to skip): " input_gh_user
GITHUB_USER="${input_gh_user:-}"

if [ -n "$GITHUB_USER" ]; then
    read -rp "GitHub PAT: " input_pat
    GITHUB_PAT="${input_pat:-}"

    if [ -n "$GITHUB_PAT" ]; then
        # Validate PAT format (should start with ghp_ or github_pat_)
        if ! [[ "$GITHUB_PAT" =~ ^(ghp_|github_pat_) ]]; then
            warn "PAT doesn't look right — expected it to start with ghp_ or github_pat_"
            read -p "Continue anyway? (y/n) " -n 1 -r
            echo ""
            if [[ ! $REPLY =~ ^[Yy]$ ]]; then
                GITHUB_PAT=""
                GITHUB_USER=""
                warn "Skipping GitHub setup — you can configure it later"
            fi
        fi
    else
        warn "No PAT provided — skipping GitHub setup"
        GITHUB_USER=""
    fi
else
    GITHUB_PAT=""
fi

# Confirm
echo ""
echo -e "${CYAN}━━━ Confirm settings ━━━${NC}"
echo ""
echo "  SSH alias:      $SSH_ALIAS"
echo "  Deploy user:    $DEPLOY_USER"
echo "  Git name:       $GIT_NAME"
echo "  Git email:      $GIT_EMAIL"
echo "  GitHub user:    $([ -n "$GITHUB_USER" ] && echo "$GITHUB_USER" || echo 'Not set')"
echo "  Git PAT:        $([ -n "$GITHUB_PAT" ] && echo 'Set' || echo 'Not set')"
echo "  Timezone:       $TIMEZONE"
echo ""
read -p "Continue? (y/n) " -n 1 -r
echo ""
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "Aborted."
    exit 0
fi

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

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

pacman -Syu --noconfirm
pacman -S --noconfirm --needed \
    ufw \
    fail2ban \
    curl \
    tmux \
    git \
    docker \
    docker-compose

# Docker must be running before ONCE installs (the ONCE installer does not
# install Docker on Arch — only on apt-based distros)
systemctl enable --now docker.service

log "System updated, packages installed, Docker running"

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

step "2/8 · Timezone & locale"

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/8 · 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

# Allow sudo without password for deploy user
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"

# 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/8 · 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/8 · Firewall (UFW)"

ufw default deny incoming
ufw default allow outgoing
ufw allow 22/tcp comment 'SSH'
ufw allow 80/tcp comment 'HTTP'
ufw allow 443/tcp comment 'HTTPS'

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

log "UFW enabled — ports 22, 80, 443 open"
ufw status verbose

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

step "6/8 · 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/8 · 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. ONCE (Docker already installed via pacman in step 1)
# =============================================================================

step "8/8 · ONCE"

log "Installing ONCE..."
su - $DEPLOY_USER -c 'curl https://get.once.com | ONCE_INTERACTIVE=false sh' || {
    warn "ONCE auto-install had issues — you can install manually:"
    warn "  ssh $SSH_ALIAS"
    warn "  curl https://get.once.com | sh"
}

# Ensure deploy user is in docker group
if getent group docker &>/dev/null; then
    usermod -aG docker $DEPLOY_USER
    log "$DEPLOY_USER added to docker group"
fi

# =============================================================================
# SWAP (2GB)
# =============================================================================

step "Swap file"

if [ ! -f /swapfile ]; then
    fallocate -l 2G /swapfile
    chmod 600 /swapfile
    mkswap /swapfile
    swapon /swapfile
    grep -q '/swapfile' /etc/fstab || echo '/swapfile none swap sw 0 0' >> /etc/fstab
    log "2GB swap created and enabled"
else
    warn "Swap file already exists — skipping"
fi

# =============================================================================
# DISABLE UNNECESSARY SERVICES
# =============================================================================

step "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)
# =============================================================================

step "Cap journal logs"

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"

# =============================================================================
# CLEAN PACMAN CACHE (keep last 2 versions of each package)
# =============================================================================

step "Pacman cache hygiene"

pacman -S --noconfirm --needed pacman-contrib
systemctl enable --now paccache.timer
log "paccache.timer enabled — old package versions pruned weekly"

# =============================================================================
# GIT CONFIGURATION
# =============================================================================

step "Git configuration"

su - $DEPLOY_USER << GITEOF
git config --global user.name "$GIT_NAME"
git config --global user.email "$GIT_EMAIL"
git config --global credential.helper store
git config --global init.defaultBranch main
GITEOF

log "Git configured for $DEPLOY_USER"

# Store PAT if provided (correct format: https://USERNAME:TOKEN@github.com)
if [ -n "$GITHUB_PAT" ] && [ -n "$GITHUB_USER" ]; then
    su - $DEPLOY_USER -c "echo 'https://$GITHUB_USER:$GITHUB_PAT@github.com' > /home/$DEPLOY_USER/.git-credentials"
    chmod 600 /home/$DEPLOY_USER/.git-credentials
    chown $DEPLOY_USER:$DEPLOY_USER /home/$DEPLOY_USER/.git-credentials
    log "GitHub PAT stored for git (https://$GITHUB_USER:***@github.com)"

    # Log Docker into GHCR so Once can pull container images
    log "Logging Docker into ghcr.io..."
    if echo "$GITHUB_PAT" | su - $DEPLOY_USER -c "docker login ghcr.io -u $GITHUB_USER --password-stdin"; then
        log "Docker logged into ghcr.io as $GITHUB_USER"
    else
        warn "Docker GHCR login failed — you can run manually: docker login ghcr.io -u $GITHUB_USER"
    fi
fi

# Clone repos if specified
if [ -n "$REPOS_TO_CLONE" ]; then
    IFS=',' read -ra REPO_PAIRS <<< "$REPOS_TO_CLONE"
    for pair in "${REPO_PAIRS[@]}"; do
        IFS=':' read -r repo dir <<< "$pair"
        if [ ! -d "$dir" ]; then
            mkdir -p "$(dirname "$dir")"
            su - $DEPLOY_USER -c "git clone https://github.com/$repo.git $dir"
            chown -R $DEPLOY_USER:$DEPLOY_USER "$dir"
            log "Cloned $repo → $dir"
        else
            warn "$dir already exists — skipping clone"
        fi
    done
fi

# =============================================================================
# CREATE DIRECTORIES
# =============================================================================

step "Directory structure"

mkdir -p /opt/apps
chown -R $DEPLOY_USER:$DEPLOY_USER /opt/apps
log "Created /opt/apps for custom Docker apps"

mkdir -p /var/www
chown -R $DEPLOY_USER:$DEPLOY_USER /var/www
log "Created /var/www for static sites"

# =============================================================================
# DEPLOY HELPER SCRIPT
# =============================================================================

step "Deploy helper script"

cat > /home/$DEPLOY_USER/deploy.sh << 'DEPLOYSCRIPT'
#!/bin/bash
# Usage: ./deploy.sh <app-directory> [branch]
set -euo pipefail

APP_DIR="${1:?Usage: ./deploy.sh <app-directory> [branch]}"
BRANCH="${2:-main}"

if [ ! -d "$APP_DIR" ]; then
    echo "Error: $APP_DIR does not exist"
    exit 1
fi

cd "$APP_DIR"
echo "📦 Deploying $(basename $APP_DIR) from branch $BRANCH..."

echo "⬇️  Pulling latest..."
git fetch origin
git checkout "$BRANCH"
git pull origin "$BRANCH"

if [ -f "docker-compose.yml" ] || [ -f "compose.yml" ]; then
    echo "🐳 Docker Compose detected — rebuilding..."
    docker compose down
    docker compose up -d --build
    echo "✅ Containers running:"
    docker compose ps
elif [ -f "Dockerfile" ]; then
    APP_NAME=$(basename "$APP_DIR")
    echo "🐳 Dockerfile detected — rebuilding $APP_NAME..."
    docker build -t "$APP_NAME" .
    docker stop "$APP_NAME" 2>/dev/null || true
    docker rm "$APP_NAME" 2>/dev/null || true
    docker run -d --name "$APP_NAME" --restart unless-stopped "$APP_NAME"
    echo "✅ Container running"
else
    echo "📁 Static site — no build step needed"
fi

echo "🎉 Deployed at $(date)"
DEPLOYSCRIPT

chown $DEPLOY_USER:$DEPLOY_USER /home/$DEPLOY_USER/deploy.sh
chmod +x /home/$DEPLOY_USER/deploy.sh
log "Deploy helper created at ~/deploy.sh"

# =============================================================================
# GENERATE LOCAL SSH CONFIG FILE
# =============================================================================

step "SSH config for your local machine"

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")

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
log "SSH config snippet saved to ~/ssh-config-snippet.txt"

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

echo ""
echo -e "${CYAN}"
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║                    SETUP 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 "  ├─ 80/tcp          HTTP"
echo "  └─ 443/tcp         HTTPS"
echo ""
echo -e "  ${GREEN}Services${NC}"
echo "  ├─ Docker:        $(docker --version 2>/dev/null || echo 'installed')"
echo "  ├─ ONCE:          run 'once' as $DEPLOY_USER to manage apps"
echo "  ├─ Fail2ban:      active (systemd backend)"
echo "  └─ Auto-updates:  weekly pacman-update.timer"
echo ""
echo -e "  ${GREEN}Paths${NC}"
echo "  ├─ Custom apps:   /opt/apps/"
echo "  ├─ Static sites:  /var/www/"
echo "  ├─ Deploy script: ~/deploy.sh <app-dir> [branch]"
echo "  └─ SSH config:    ~/ssh-config-snippet.txt"
echo ""
echo -e "  ${YELLOW}NEXT STEPS${NC}"
echo ""
echo "  1. On your local machine, run this to 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. Test SSH:  ssh $SSH_ALIAS"
echo "  3. Run ONCE:  ssh $SSH_ALIAS 'once'"
echo "  4. Deploy:    ssh $SSH_ALIAS '~/deploy.sh /opt/apps/myapp'"
echo ""
echo "  Once step 1 is done, Claude Code can control this server."
echo ""
echo -e "  ${YELLOW}A reboot is recommended to apply kernel updates:${NC}"
echo -e "  ${CYAN}  reboot${NC}"
echo ""