#!/bin/bash
set -e

# =============================================================================
# machine-setup — provisioning for a fresh machine
#
# Brings a blank box up to a usable state: users, SSH, networking, firewall,
# Docker, shell and editor tooling, language runtimes.
#
# Run as root:  sudo scripts/setup/machine-setup/machine-setup.sh
# =============================================================================

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROGRESS_FILE="$SCRIPT_DIR/.setup-progress"

# Shared state, output helpers, the step/resume machine and OS detection. Kept in
# lib/ so a step can eventually be read — or run — on its own without dragging the
# whole script in. Definitions only; nothing in there acts.
# shellcheck source=lib/base.sh
source "$SCRIPT_DIR/lib/base.sh"
# shellcheck source=lib/packages.sh
source "$SCRIPT_DIR/lib/packages.sh"

# Trap errors with context. Installed here rather than in lib/base.sh, because
# that file is definitions only and a trap is a side effect on whoever sources it.
trap 'echo ""; echo -e "${RED}╔══════════════════════════════════════════════════╗${NC}"; echo -e "${RED}║  SETUP FAILED${NC}"; echo -e "${RED}║  Step: ${CURRENT_STEP:-unknown}${NC}"; echo -e "${RED}║  Line: $LINENO${NC}"; echo -e "${RED}║  Command: $BASH_COMMAND${NC}"; echo -e "${RED}╚══════════════════════════════════════════════════╝${NC}"' ERR

# =============================================================================
# 1. Pre-flight
# =============================================================================

echo ""
echo -e "${BOLD}╔══════════════════════════════════════════════════╗${NC}"
echo -e "${BOLD}║          Machine Setup                           ║${NC}"
echo -e "${BOLD}╚══════════════════════════════════════════════════╝${NC}"

detect_os
echo ""
info "Machine:  ${OS_NAME} (${ARCH})"
info "Packages: ${PM:-none detected}"
[[ "$IS_WSL" == true ]] && warn "WSL detected — the suspend, logind and boot-hang steps do not apply here"

# Everything below this line is written against apt and systemd. Detection above
# recognises pacman, dnf and brew so the branches have somewhere to hang, but
# nothing implements them yet — and running the apt path on Arch would half-build
# a machine and stop somewhere unhelpful. Refuse clearly instead, and relax this
# list one entry at a time as each package manager grows a real path.
case "$PM" in
  apt) ;;
  "") fail "Could not find a package manager for '${OS_NAME}'." ;;
  *) fail "${OS_NAME} uses ${PM}, which this script does not implement yet — apt-based systems only, so far." ;;
esac

ask_machine_role
info "Role:     ${MACHINE_ROLE}"

if [[ -f "$PROGRESS_FILE" ]]; then
  DONE_COUNT=$(wc -l < "$PROGRESS_FILE")
  echo -e "${YELLOW}  Resuming — $DONE_COUNT step(s) already completed${NC}"
  echo -e "${YELLOW}  Progress file: $PROGRESS_FILE${NC}"
  echo -e "${YELLOW}  To start fresh: rm $PROGRESS_FILE${NC}"
fi

if [[ "$EUID" -ne 0 ]]; then
  fail "Please run as root: sudo ./machine-setup.sh"
fi

prompt_value USERNAME "New admin username (or existing)" ""
if [[ -z "$USERNAME" ]]; then
  fail "Username cannot be empty"
fi

USER_HOME="/home/$USERNAME"

# =============================================================================
# 2. System Update & Essentials
# =============================================================================

step "System Update & Essentials"
if ! skip; then
  info "Refreshing the package index..."
  pkg_refresh

  # The one place that deliberately moves versions of things already installed.
  # Everything else in this script only ever adds what is absent.
  info "Upgrading installed packages..."
  pkg_upgrade_all

  # shellcheck disable=SC2046  # word splitting is how the list is passed
  pkg_install "Core packages" $(pkgs_core)

  ok "System updated and essentials installed"

  # lazydocker
  info "Installing lazydocker..."
  curl -fsSL https://raw.githubusercontent.com/jesseduffield/lazydocker/master/scripts/install_update_linux.sh | DIR=/usr/local/bin bash
  ok "lazydocker installed"

  # lazygit
  info "Installing lazygit..."
  LAZYGIT_VERSION=$(curl -fsSL "https://api.github.com/repos/jesseduffield/lazygit/releases/latest" | jq -r '.tag_name' | tr -d 'v')
  curl -fsSLo /tmp/lazygit.tar.gz "https://github.com/jesseduffield/lazygit/releases/download/v${LAZYGIT_VERSION}/lazygit_${LAZYGIT_VERSION}_Linux_x86_64.tar.gz"
  tar -C /usr/local/bin -xzf /tmp/lazygit.tar.gz lazygit
  rm -f /tmp/lazygit.tar.gz
  ok "lazygit installed"

  # starship
  info "Installing starship..."
  curl -fsSL https://starship.rs/install.sh | sh -s -- -y
  ok "starship installed"

  # fastfetch
  info "Installing fastfetch..."
  add-apt-repository -y ppa:zhangsongcui3371/fastfetch > /dev/null 2>&1 || true
  apt-get update -y > /dev/null 2>&1
  apt-get install -y fastfetch
  ok "fastfetch installed"

  SUMMARY+=("System packages updated and essentials installed")
  step_ok
fi

# ─────────────────────────────────────────────────────────────────────────────
# WIP BOUNDARY — everything below is still the original script
# ─────────────────────────────────────────────────────────────────────────────
#
# Move this block down as each section is worked through, and delete it when the
# walk reaches the end. It exists so the rewritten part can be run start to
# finish on its own without the untouched sections acting on the machine.
echo ""
info "Stopping here — sections below this point have not been rewritten yet."
info "Run again from scratch with:  sudo rm -f $PROGRESS_FILE"
echo ""
exit 0

# =============================================================================
# 3. Locale
# =============================================================================

step "Locale"
if ! skip; then
  info "Setting locale to en_US.UTF-8..."
  locale-gen en_US.UTF-8 || true
  update-locale LANG=en_US.UTF-8

  ok "Locale set"
  SUMMARY+=("Locale: en_US.UTF-8")
  step_ok
fi

# =============================================================================
# 4. Timezone
# =============================================================================

step "Timezone"
if ! skip; then
  info "Select timezone:"
  TZ_OPTIONS=("UTC" "Europe/Lisbon" "Europe/Stockholm" "Europe/London" "US/Eastern" "US/Pacific" "Asia/Tokyo")
  for i in "${!TZ_OPTIONS[@]}"; do
    echo "  [$((i+1))] ${TZ_OPTIONS[$i]}"
  done

  prompt_value TZ_CHOICE "Pick a number or type a timezone" "1"

  # If the input is a number, map to the array; otherwise use as-is
  if [[ "$TZ_CHOICE" =~ ^[0-9]+$ ]] && (( TZ_CHOICE >= 1 && TZ_CHOICE <= ${#TZ_OPTIONS[@]} )); then
    TIMEZONE="${TZ_OPTIONS[$((TZ_CHOICE-1))]}"
  else
    TIMEZONE="$TZ_CHOICE"
  fi

  timedatectl set-timezone "$TIMEZONE"
  ok "Timezone: $TIMEZONE"
  SUMMARY+=("Timezone: $TIMEZONE")
  step_ok
fi

# =============================================================================
# 5. Swap
# =============================================================================

step "Swap"
if ! skip; then
  if ! swapon --show | grep -q '/'; then
    info "Creating swap file..."
    RAM_KB=$(grep MemTotal /proc/meminfo | awk '{print $2}')
    if (( RAM_KB <= 2097152 )); then
      SWAP_SIZE="2G"
    elif (( RAM_KB <= 8388608 )); then
      SWAP_SIZE="4G"
    else
      SWAP_SIZE="8G"
    fi

    fallocate -l "$SWAP_SIZE" /swapfile
    chmod 600 /swapfile
    mkswap /swapfile
    swapon /swapfile

    if ! grep -q '/swapfile' /etc/fstab; then
      echo '/swapfile none swap sw 0 0' >> /etc/fstab
    fi

    ok "Swap: $SWAP_SIZE"
    SUMMARY+=("Swap: $SWAP_SIZE created")
  else
    ok "Swap already exists, skipping"
    SUMMARY+=("Swap: already present")
  fi

  # Set swappiness
  sysctl vm.swappiness=10
  if ! grep -q 'vm.swappiness' /etc/sysctl.conf; then
    echo 'vm.swappiness=10' >> /etc/sysctl.conf
  else
    sed -i 's/^vm.swappiness=.*/vm.swappiness=10/' /etc/sysctl.conf
  fi
  step_ok
fi

# =============================================================================
# 6. Auto-Suspend (disable for servers)
# =============================================================================

step "Auto-Suspend (disable for servers)"
if ! skip; then
  info "Disabling auto-suspend and idle actions..."

  # Mask sleep/suspend targets
  systemctl mask sleep.target suspend.target hibernate.target hybrid-sleep.target 2>/dev/null || true

  # Configure logind — patch individual keys, don't overwrite the file
  set_logind_key() {
    local key="$1" val="$2" file="/etc/systemd/logind.conf"
    if grep -qE "^${key}=" "$file" 2>/dev/null; then
      sed -i "s|^${key}=.*|${key}=${val}|" "$file"
    elif grep -qE "^#${key}=" "$file" 2>/dev/null; then
      sed -i "s|^#${key}=.*|${key}=${val}|" "$file"
    else
      echo "${key}=${val}" >> "$file"
    fi
  }
  set_logind_key HandleLidSwitch ignore
  set_logind_key HandleLidSwitchExternalPower ignore
  set_logind_key HandlePowerKey ignore
  set_logind_key IdleAction none
  set_logind_key RuntimeDirectorySize 10%

  systemctl restart systemd-logind

  ok "Auto-suspend disabled"
  SUMMARY+=("Auto-suspend: disabled")
  step_ok
fi

# =============================================================================
# 6.5 Boot Hang Fix (NetworkManager desktops only) — OPTIONAL, default off
# =============================================================================
#
# On Ubuntu Desktop, NetworkManager owns the NIC and systemd-networkd runs
# nothing — so systemd-networkd-wait-online.service blocks boot for its full
# ~120s timeout. Masking it removes the hang.
#
# On Ubuntu Server / cloud VPS images, systemd-networkd IS the network manager
# (cloud-init + netplan). There the service is load-bearing and completes
# quickly, so there's nothing to fix. The guard below detects which stack is
# live and only acts on the NetworkManager case.
#
# NOTE: do NOT `systemctl disable --now systemd-networkd.{socket,service}` here.
# On a networkd-managed VPS that brings the box up with no network on the next
# boot — no SSH, provider rescue console only. Keep it as a manual, desktop-only
# step if you ever want it.

step "Boot Hang Fix"
if ! skip; then
  prompt_value FIX_BOOT_HANG "Mask systemd-networkd-wait-online on NetworkManager boxes? (y/n)" "n"

  if [[ "$FIX_BOOT_HANG" != "y" ]]; then
    info "Skipping boot hang fix"
    SUMMARY+=("Boot hang fix: skipped")
  else
    # Only mask the spurious wait-online where NetworkManager is active and
    # networkd is not. Anything else (networkd stack, or unclear) is left alone.
    if systemctl is-active --quiet NetworkManager.service \
       && ! systemctl is-active --quiet systemd-networkd.service; then
      systemctl mask --now systemd-networkd-wait-online.service
      ok "Masked systemd-networkd-wait-online (NetworkManager stack)"
      SUMMARY+=("Boot hang fix: wait-online masked (NetworkManager)")
    else
      warn "networkd stack (or unclear) — leaving wait-online alone"
      SUMMARY+=("Boot hang fix: not applied (networkd stack)")
    fi
  fi
  step_ok
fi

# =============================================================================
# 7. User Creation
# =============================================================================

step "User Creation"
if ! skip; then
  if id "$USERNAME" &>/dev/null; then
    info "User '$USERNAME' already exists, skipping creation"
    SUMMARY+=("User: $USERNAME (existing)")
  else
    info "Creating user '$USERNAME'..."
    adduser --gecos "" "$USERNAME"
    usermod -aG sudo "$USERNAME"
    ok "User '$USERNAME' created with sudo access"
    SUMMARY+=("User: $USERNAME created")
  fi

  prompt_value PASSWORDLESS_SUDO "Enable passwordless sudo for $USERNAME? (y/n)" "y"

  if [[ "$PASSWORDLESS_SUDO" == "y" ]]; then
    echo "$USERNAME ALL=(ALL) NOPASSWD: ALL" > "/etc/sudoers.d/99-$USERNAME-nopasswd"
    chmod 440 "/etc/sudoers.d/99-$USERNAME-nopasswd"
    visudo -cf "/etc/sudoers.d/99-$USERNAME-nopasswd" || rm -f "/etc/sudoers.d/99-$USERNAME-nopasswd"
    ok "Passwordless sudo enabled (remove /etc/sudoers.d/99-$USERNAME-nopasswd to disable)"
    SUMMARY+=("Sudo: passwordless")
  fi

  # Copy .tmux.conf if present next to this script
  if [[ -f "$SCRIPT_DIR/.tmux.conf" ]]; then
    cp "$SCRIPT_DIR/.tmux.conf" "$USER_HOME/.tmux.conf"
    chown "$USERNAME:$USERNAME" "$USER_HOME/.tmux.conf"
    ok "Copied .tmux.conf to $USER_HOME"
  fi
  step_ok
fi

# =============================================================================
# 8. SSH Keys
# =============================================================================

step "SSH Keys"
if ! skip; then
  info "SSH key setup for '$USERNAME':"
  echo "  [1] Extract from ssh-keys.zip (next to this script)"
  echo "  [2] Paste key content"
  echo "  [3] Generate new ed25519 key"

  prompt_value SSH_CHOICE "Pick an option" "1"

  SSH_DIR="$USER_HOME/.ssh"
  mkdir -p "$SSH_DIR"

  case "$SSH_CHOICE" in
    1)
      ZIP_FILE="$SCRIPT_DIR/ssh-keys.zip"
      if [[ ! -f "$ZIP_FILE" ]]; then
        warn "ssh-keys.zip not found at $ZIP_FILE"
        warn "Skipping SSH key setup"
      else
        read -rsp "Zip password: " ZIP_PASS
        echo
        if [[ -n "$ZIP_PASS" ]]; then
          unzip -o -P "$ZIP_PASS" "$ZIP_FILE" -d "$SSH_DIR"
        else
          unzip -o "$ZIP_FILE" -d "$SSH_DIR"
        fi
        ok "SSH keys extracted from zip"
        SUMMARY+=("SSH keys: extracted from zip")
      fi
      ;;
    2)
      echo "Paste your PUBLIC key (then press Enter, then Ctrl-D):"
      cat > "$SSH_DIR/id_ed25519.pub"
      echo "Paste your PRIVATE key (then press Enter, then Ctrl-D):"
      cat > "$SSH_DIR/id_ed25519"
      # Also add public key to authorized_keys
      cat "$SSH_DIR/id_ed25519.pub" >> "$SSH_DIR/authorized_keys"
      ok "SSH keys written from input"
      SUMMARY+=("SSH keys: pasted manually")
      ;;
    3)
      prompt_value SSH_EMAIL "Email/comment for key" ""
      sudo -u "$USERNAME" ssh-keygen -t ed25519 -C "$SSH_EMAIL" -f "$SSH_DIR/id_ed25519" -N ""
      cat "$SSH_DIR/id_ed25519.pub" >> "$SSH_DIR/authorized_keys"
      ok "SSH key generated"
      echo "   Public key:"
      cat "$SSH_DIR/id_ed25519.pub"
      SUMMARY+=("SSH keys: generated new ed25519")
      ;;
  esac

  chmod 700 "$SSH_DIR"
  find "$SSH_DIR" -type f -exec chmod 600 {} \;
  chown -R "$USERNAME:$USERNAME" "$SSH_DIR"
  step_ok
fi

# =============================================================================
# 9. SSH Hardening
# =============================================================================

step "SSH Hardening"
if ! skip; then
  info "Hardening SSH..."
  sed -i 's/^#\?PasswordAuthentication .*/PasswordAuthentication no/' /etc/ssh/sshd_config
  sed -i 's/^#\?PermitRootLogin .*/PermitRootLogin no/' /etc/ssh/sshd_config
  sed -i 's/^#\?ChallengeResponseAuthentication .*/ChallengeResponseAuthentication no/' /etc/ssh/sshd_config
  systemctl restart ssh

  ok "SSH hardened (password auth disabled, root login disabled)"
  SUMMARY+=("SSH: hardened")
  step_ok
fi

# =============================================================================
# 10. DNS
# =============================================================================

step "DNS"
if ! skip; then
  info "Configuring DNS (Cloudflare + Google)..."
  PRIMARY_DNS=("1.1.1.1" "8.8.8.8")
  FALLBACK_DNS=("1.0.0.1" "8.8.4.4")

  if systemctl is-active --quiet systemd-resolved; then
    cp /etc/systemd/resolved.conf "/etc/systemd/resolved.conf.bak.$(date +%s)" 2>/dev/null || true

    cat > /etc/systemd/resolved.conf <<EOF
[Resolve]
DNS=${PRIMARY_DNS[*]}
FallbackDNS=${FALLBACK_DNS[*]}
EOF

    # Ensure resolv.conf points to the stub
    if [[ -L /etc/resolv.conf ]]; then
      TARGET=$(readlink -f /etc/resolv.conf || true)
      if [[ "$TARGET" != "/run/systemd/resolve/stub-resolv.conf" && "$TARGET" != "/run/systemd/resolve/resolv.conf" ]]; then
        ln -sf /run/systemd/resolve/stub-resolv.conf /etc/resolv.conf
      fi
    else
      cp /etc/resolv.conf "/etc/resolv.conf.bak.$(date +%s)"
      ln -sf /run/systemd/resolve/stub-resolv.conf /etc/resolv.conf
    fi

    systemctl restart systemd-resolved
    ok "DNS configured (systemd-resolved)"
  else
    if [[ -e /etc/resolv.conf ]] || [[ -L /etc/resolv.conf ]]; then
      cp -a /etc/resolv.conf "/etc/resolv.conf.bak.$(date +%s)" 2>/dev/null || true
    fi
    if lsattr /etc/resolv.conf 2>/dev/null | grep -q 'i-'; then
      chattr -i /etc/resolv.conf
    fi

    {
      echo "# Generated by setup-ubuntu.sh on $(date -u +'%Y-%m-%dT%H:%M:%SZ')"
      for ns in "${PRIMARY_DNS[@]}"; do
        echo "nameserver $ns"
      done
    } > /etc/resolv.conf

    ok "DNS configured (static resolv.conf)"
  fi

  SUMMARY+=("DNS: Cloudflare + Google")

  # Test DNS
  if command -v dig &>/dev/null; then
    dig +short google.com > /dev/null 2>&1 && ok "DNS resolution works" || warn "DNS test failed"
  else
    getent hosts google.com > /dev/null 2>&1 && ok "DNS resolution works" || warn "DNS test failed"
  fi
  step_ok
fi

# =============================================================================
# 11. Static IP
# =============================================================================

step "Static IP"
if ! skip; then
  # Detect primary interface and current IP
  IFACE=$(ip route get 8.8.8.8 2>/dev/null | awk '{for(i=1;i<=NF;i++) if($i=="dev"){print $(i+1); exit}}')
  CURRENT_IP=$(ip -4 addr show "$IFACE" 2>/dev/null | grep -oP '(?<=inet\s)\d+(\.\d+){3}/\d+' | head -1)
  GATEWAY=$(ip route | grep default | awk '{print $3}' | head -1)

  if [[ -z "$IFACE" || -z "$CURRENT_IP" || -z "$GATEWAY" ]]; then
    warn "Could not detect network interface, IP, or gateway — skipping static IP"
    SUMMARY+=("Static IP: skipped (detection failed)")
  else
    # Check if a static Netplan config already exists for this interface
    EXISTING_STATIC=$(grep -rl "dhcp4: false\|dhcp4: no" /etc/netplan/ 2>/dev/null | head -1)
    if [[ -n "$EXISTING_STATIC" ]]; then
      warn "Static IP already configured in $EXISTING_STATIC — skipping"
      SUMMARY+=("Static IP: already static (skipped)")
    else
      info "Detected network configuration:"
      echo "    Interface : $IFACE"
      echo "    IP        : $CURRENT_IP"
      echo "    Gateway   : $GATEWAY"
      echo ""
      echo -e "${RED}${BOLD}╔══════════════════════════════════════════════════╗${NC}"
      echo -e "${RED}${BOLD}║  STOP — DO NOT DO THIS ON A VPS / CLOUD SERVER   ║${NC}"
      echo -e "${RED}${BOLD}╚══════════════════════════════════════════════════╝${NC}"
      echo -e "${YELLOW}  This freezes the CURRENT DHCP lease into netplan and applies"
      echo -e "  it immediately. On a VPS or cloud instance that can strand the box:${NC}"
      echo ""
      echo -e "    ${RED}!${NC} Providers hand out addresses by DHCP and may reassign them —"
      echo -e "      the hardcoded IP then belongs to someone else."
      echo -e "    ${RED}!${NC} cloud-init regenerates netplan config on boot; the two configs"
      echo -e "      can conflict and leave the interface with no usable address."
      echo -e "    ${RED}!${NC} The gateway/netmask on cloud networks is often NOT what a"
      echo -e "      simple route lookup suggests (point-to-point, /32 routing)."
      echo -e "    ${RED}!${NC} If any value is wrong, ${BOLD}netplan apply kills the SSH session"
      echo -e "      you are on right now${NC} and you are locked out until the provider's"
      echo -e "      rescue console."
      echo ""
      echo -e "${YELLOW}  Only answer 'y' on a machine you physically control (LAN box,"
      echo -e "  homelab, bare metal) where the address is yours to pin — and even"
      echo -e "  then, a DHCP reservation on the router is the safer way to do it.${NC}"
      echo ""
      prompt_value CONFIRM_STATIC "Set this as static IP? (y/n)" "n"

      if [[ "$CONFIRM_STATIC" != "y" ]]; then
        info "Skipping static IP configuration"
        SUMMARY+=("Static IP: skipped (user declined)")
      else
        NETPLAN_FILE="/etc/netplan/99-static-ip.yaml"
        cat > "$NETPLAN_FILE" <<EOF
network:
  version: 2
  ethernets:
    $IFACE:
      dhcp4: false
      addresses:
        - $CURRENT_IP
      routes:
        - to: default
          via: $GATEWAY
      nameservers:
        addresses: [1.1.1.1, 8.8.8.8]
EOF
        chmod 600 "$NETPLAN_FILE"
        netplan apply
        ok "Static IP set: $CURRENT_IP on $IFACE (gateway: $GATEWAY)"
        SUMMARY+=("Static IP: $CURRENT_IP on $IFACE")
      fi
    fi
  fi
  step_ok
fi

# =============================================================================
# 12. Fail2ban
# =============================================================================

step "Fail2ban"
if ! skip; then
  info "Installing fail2ban..."
  apt-get install -y fail2ban
  systemctl enable fail2ban
  systemctl start fail2ban

  ok "Fail2ban active"
  SUMMARY+=("Fail2ban: installed and active")
  step_ok
fi

# =============================================================================
# 13. Unattended Upgrades
# =============================================================================

step "Unattended Upgrades"
if ! skip; then
  info "Configuring unattended upgrades..."
  apt-get install -y unattended-upgrades

  cat > /etc/apt/apt.conf.d/20auto-upgrades <<EOF
APT::Periodic::Update-Package-Lists "1";
APT::Periodic::Unattended-Upgrade "1";
APT::Periodic::AutocleanInterval "7";
EOF

  ok "Unattended upgrades enabled"
  SUMMARY+=("Unattended upgrades: enabled")
  step_ok
fi

# =============================================================================
# 14. Git Config
# =============================================================================

step "Git Config"
if ! skip; then
  info "Git configuration for '$USERNAME'..."
  prompt_value GIT_NAME "Git user.name" ""
  prompt_value GIT_EMAIL "Git user.email" ""

  as_user "git config --global user.name $(printf '%q' "$GIT_NAME")"
  as_user "git config --global user.email $(printf '%q' "$GIT_EMAIL")"
  as_user "git config --global init.defaultBranch main"
  as_user "git config --global core.editor nvim"

  ok "Git configured"
  SUMMARY+=("Git: $GIT_NAME <$GIT_EMAIL>")
  step_ok
fi

# =============================================================================
# 15. Docker
# =============================================================================

step "Docker"
if ! skip; then
  info "Installing Docker CE..."
  install -m 0755 -d /etc/apt/keyrings
  curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --batch --yes --dearmor -o /etc/apt/keyrings/docker.gpg
  chmod a+r /etc/apt/keyrings/docker.gpg

  echo \
    "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] \
    https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" \
    > /etc/apt/sources.list.d/docker.list

  apt-get update -y
  apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

  usermod -aG docker "$USERNAME"

  # Create shared docker network
  DOCKER_NETWORK="${SETUP_DOCKER_NETWORK:-services}"
  docker network create "$DOCKER_NETWORK" 2>/dev/null || true

  ok "Docker installed, '$USERNAME' added to docker group"
  SUMMARY+=("Docker: installed")
  step_ok
fi
# =============================================================================
# 16. Oh My Zsh + Prompt (as user)
# =============================================================================

step "Oh My Zsh + Prompt"
if ! skip; then
  info "Installing Zsh and Oh My Zsh..."

  apt-get install -y zsh

  as_user 'sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended'

  ZSHRC="$USER_HOME/.zshrc"

  cat >> "$ZSHRC" <<'EOF'

# Starship prompt
eval "$(starship init zsh)"
EOF

  chsh -s /bin/zsh "$USERNAME"

  # Starship config — disable runtime/language version modules
  mkdir -p "$USER_HOME/.config"
  cat > "$USER_HOME/.config/starship.toml" <<'STARSHIP_EOF'
add_newline = true
command_timeout = 200

[bun]
disabled = true

[nodejs]
disabled = true

[python]
disabled = true

[rust]
disabled = true

[golang]
disabled = true

[package]
disabled = true
STARSHIP_EOF
  chown -R "$USERNAME:$USERNAME" "$USER_HOME/.config"

  ok "Zsh + Oh My Zsh installed, starship prompt set"
  SUMMARY+=("Shell: zsh with Oh My Zsh + Starship")
  step_ok
fi

# =============================================================================
# 17. Tailscale (as user)
# =============================================================================

step "Tailscale"
if ! skip; then
  prompt_value INSTALL_TAILSCALE "Install Tailscale? (y/n)" "y"

  if [[ "$INSTALL_TAILSCALE" != "y" || "${SETUP_SKIP_TAILSCALE:-}" == "true" ]]; then
    info "Skipping Tailscale"
    SUMMARY+=("Tailscale: skipped")
  else
    info "Installing Tailscale..."
    curl -fsSL https://tailscale.com/install.sh | sh

    # Enable IP forwarding
    sed -i '/^net.ipv4.ip_forward/d' /etc/sysctl.conf
    sed -i '/^net.ipv6.conf.all.forwarding/d' /etc/sysctl.conf
    echo "net.ipv4.ip_forward=1" >> /etc/sysctl.conf
    echo "net.ipv6.conf.all.forwarding=1" >> /etc/sysctl.conf
    sysctl -p

    prompt_value TS_LOGIN_SERVER "Tailscale login server" "https://headscale.pastilhas.eu"
    prompt_value TS_AUTHKEY "Tailscale auth key" ""

    prompt_value TS_EXIT_NODE "Act as exit node? (y/n)" "n"
    TS_ADVERTISE=""

    if [[ "$TS_EXIT_NODE" == "y" ]]; then
      TS_ADVERTISE="--advertise-exit-node"

      cat > /etc/sysctl.d/99-tailscale-exit.conf <<EOF
net.ipv4.ip_forward = 1
net.ipv6.conf.all.forwarding = 1
EOF
      sysctl --system > /dev/null

      # NIC offload optimizations via networkd-dispatcher
      if ! command -v networkd-dispatcher &>/dev/null; then
        apt-get install -y networkd-dispatcher
      fi
      mkdir -p /etc/networkd-dispatcher/routable.d

      cat > /etc/networkd-dispatcher/routable.d/50-tailscale-exit <<'EXITSCRIPT'
#!/usr/bin/env bash
set -Eeuo pipefail

IF="${IFACE:-}"
if [[ -z "${IF}" ]]; then
  IF="$(ip -o route get 8.8.8.8 2>/dev/null | awk '{for(i=1;i<=NF;i++) if($i=="dev"){print $(i+1); exit}}')"
fi

[[ -n "${IF}" ]] || exit 0
command -v ethtool >/dev/null 2>&1 || exit 0

ethtool -k "${IF}" 2>/dev/null | grep -q "^generic-receive-offload: " && ethtool -K "${IF}" gro on || true
ethtool -k "${IF}" 2>/dev/null | grep -q "^rx-udp-gro-forwarding: " && ethtool -K "${IF}" rx-udp-gro-forwarding on || true
ethtool -k "${IF}" 2>/dev/null | grep -q "^large-receive-offload: " && ethtool -K "${IF}" lro off || true

exit 0
EXITSCRIPT

      chmod 755 /etc/networkd-dispatcher/routable.d/50-tailscale-exit
      systemctl enable --now networkd-dispatcher > /dev/null 2>&1 || true

      # Run once now
      CURRENT_IF="$(ip -o route get 8.8.8.8 | awk '{for(i=1;i<=NF;i++) if($i=="dev"){print $(i+1); exit}}')"
      IFACE="$CURRENT_IF" bash /etc/networkd-dispatcher/routable.d/50-tailscale-exit || true
      ok "Exit node NIC optimizations applied on: ${CURRENT_IF:-unknown}"
    fi

    tailscale up --login-server "$TS_LOGIN_SERVER" --authkey "$TS_AUTHKEY" ${TS_ADVERTISE:+"$TS_ADVERTISE"}

    TS_IP=$(tailscale ip -4 2>/dev/null || echo "unknown")
    ok "Tailscale connected (IP: $TS_IP)"
    SUMMARY+=("Tailscale: connected ($TS_IP)")
  fi
  step_ok
fi

# =============================================================================
# 18. Neovim + LazyVim (as user)
# =============================================================================

step "Neovim + LazyVim"
if ! skip; then
  ZSHRC="$USER_HOME/.zshrc"

  info "Installing Neovim..."
  NVIM_ARCH=$(uname -m)
  case $NVIM_ARCH in
    aarch64) NVIM_ARCH=aarch64 ;;
    *)       NVIM_ARCH=x86_64 ;;
  esac
  curl -LO "https://github.com/neovim/neovim/releases/latest/download/nvim-linux-${NVIM_ARCH}.tar.gz"
  rm -rf "/opt/nvim-linux-${NVIM_ARCH}"
  tar -C /opt -xzf "nvim-linux-${NVIM_ARCH}.tar.gz"
  rm -f "nvim-linux-${NVIM_ARCH}.tar.gz"
  ln -sf "/opt/nvim-linux-${NVIM_ARCH}/bin/nvim" /usr/local/bin/nvim

  # Add to PATH in .zshrc
  cat >> "$ZSHRC" <<EOF

# Neovim
export PATH="\$PATH:/opt/nvim-linux-${NVIM_ARCH}/bin"
EOF

  info "Choose Neovim configuration:"
  echo "  [1] Default LazyVim starter"
  echo "  [2] Custom repo"

  prompt_value NVIM_CHOICE "Pick an option" "1"

  NVIM_CONFIG="$USER_HOME/.config/nvim"

  if [[ "$NVIM_CHOICE" == "2" ]]; then
    prompt_value NVIM_REPO "Git repo for nvim config" "git@gogs:andrepadez/nvim-config.git"
  else
    NVIM_REPO="https://github.com/LazyVim/starter"
  fi

  if [[ -d "$NVIM_CONFIG" ]]; then
    warn "$NVIM_CONFIG already exists, backing up to ${NVIM_CONFIG}.bak"
    mv "$NVIM_CONFIG" "${NVIM_CONFIG}.bak"
  fi

  sudo -u "$USERNAME" -H git clone "$NVIM_REPO" "$NVIM_CONFIG"
  sudo -u "$USERNAME" rm -rf "$NVIM_CONFIG/.git" 2>/dev/null || true

  chown -R "$USERNAME:$USERNAME" "$USER_HOME/.config"
  chown "$USERNAME:$USERNAME" "$ZSHRC"

  ok "Neovim + config installed"
  SUMMARY+=("Neovim: installed with LazyVim")
  step_ok
fi

# =============================================================================
# 19. JS/TS Runtimes (system-wide for servers, optional nvm for dev)
# =============================================================================

step "JS/TS Runtimes"
if ! skip; then
  ZSHRC="$USER_HOME/.zshrc"

  if [[ "${SETUP_SKIP_RUNTIMES:-}" == "true" ]]; then
    info "Skipping JS/TS runtimes (SETUP_SKIP_RUNTIMES=true)"
    SUMMARY+=("Runtimes: skipped")
  else
    info "Installing JS/TS runtimes..."

    # Node 22 — always system-wide via NodeSource
    info "Installing system-wide Node 22..."
    curl -fsSL https://deb.nodesource.com/setup_22.x | bash -
    apt-get install -y nodejs
    npm config set prefix /usr/local
    ok "System Node 22 installed ($(node -v))"
    SUMMARY+=("Node: system-wide v22 via NodeSource")

    # Bun (user-local install)
    as_user 'curl -fsSL https://bun.sh/install | bash'
    cat >> "$ZSHRC" <<'EOF'

# Bun
export PATH="$HOME/.bun/bin:$PATH"
EOF
    ok "Bun installed"

    # pm2 (system-wide, via system npm)
    npm install -g pm2
    ok "pm2 installed (system-wide)"

    # Deno (user-local install)
    as_user 'curl -fsSL https://deno.land/install.sh | sh'
    cat >> "$ZSHRC" <<'EOF'

# Deno
export DENO_INSTALL="$HOME/.deno"
export PATH="$DENO_INSTALL/bin:$PATH"
EOF
    ok "Deno installed"

    chown "$USERNAME:$USERNAME" "$ZSHRC"
  fi
  step_ok
fi

# =============================================================================
# 19.5 Dev Tools (system-wide for servers, local for dev users)
# =============================================================================

step "Dev Tools"
if ! skip; then
  if [[ "${SETUP_SKIP_DEVTOOLS:-}" == "true" ]]; then
    info "Skipping dev tools (SETUP_SKIP_DEVTOOLS=true)"
    SUMMARY+=("Dev tools: skipped")
  else
    info "Select dev tools to install:"
    echo "  [1] Claude Code"
    echo "  [2] Opencode"
    echo "  [3] PI"

    prompt_value DEV_TOOLS "Enter numbers separated by spaces, or 'all'" "all"

    if [[ "$DEV_TOOLS" == "all" ]]; then
      DEV_TOOLS="1 2 3"
    fi

    DEV_INSTALLED=()

    for tool in $DEV_TOOLS; do
      case "$tool" in
        1)
          info "Installing Claude Code..."
          npm install -g @anthropic-ai/claude-code
          ok "Claude Code installed (system-wide)"
          DEV_INSTALLED+=("Claude Code")
          ;;
        2)
          info "Installing Opencode..."
          as_user 'curl -fsSL https://opencode.ai/install | bash'
          ok "Opencode installed"
          DEV_INSTALLED+=("Opencode")
          ;;
        3)
          info "Installing PI..."
          npm install -g @mariozechner/pi-coding-agent
          ok "PI installed (system-wide)"
          DEV_INSTALLED+=("PI")
          ;;
      esac
    done

    if [[ ${#DEV_INSTALLED[@]} -gt 0 ]]; then
      ok "Dev tools: ${DEV_INSTALLED[*]}"
      SUMMARY+=("Dev tools: ${DEV_INSTALLED[*]}")
    else
      SUMMARY+=("Dev tools: none selected")
    fi
  fi
  step_ok
fi

# =============================================================================
# 20. UFW (last, after everything)
# =============================================================================

step "UFW Firewall"
if ! skip; then
  info "Configuring UFW firewall..."

  ufw default deny incoming
  ufw default allow outgoing
  ufw allow OpenSSH

  # Allow all traffic on Tailscale interface
  ufw allow in on tailscale0

  # Docker + UFW fix
  UFW_RULES="/etc/ufw/after.rules"
  DOCKER_RULES="$SCRIPT_DIR/ufw-docker-rules.conf"

  if [[ -f "$DOCKER_RULES" ]]; then
    if ! grep -q 'DOCKER-USER' "$UFW_RULES" 2>/dev/null; then
      info "Appending Docker UFW rules..."
      echo "" >> "$UFW_RULES"
      cat "$DOCKER_RULES" >> "$UFW_RULES"
      ok "Docker UFW rules added"
    else
      ok "Docker UFW rules already present"
    fi
  else
    warn "ufw-docker-rules.conf not found, skipping Docker UFW fix"
  fi

  ufw --force enable
  ok "UFW enabled"
  SUMMARY+=("UFW: enabled (SSH, Tailscale, Docker rules)")
  step_ok
fi

# =============================================================================
# 21. Final .zshrc additions
# =============================================================================

step "Final .zshrc additions"
if ! skip; then
  ZSHRC="$USER_HOME/.zshrc"
  cat >> "$ZSHRC" <<'EOF'

# Neovim aliases
alias n='nvim'
alias vim='nvim'

# Default editor
export EDITOR='nvim'
export VISUAL='nvim'
export SUDO_EDITOR='nvim'

# Quick reload
alias sz="source ~/.zshrc"
EOF
  chown "$USERNAME:$USERNAME" "$ZSHRC"
  step_ok
fi

# =============================================================================
# 22. Emergency Disk Ballast
# =============================================================================
#
# Pre-allocates a junk file sized at 10% of currently-available disk. It holds
# no data — its only job is to be deleted when the filesystem is about to fill,
# buying enough headroom to SSH in and clean up properly instead of hitting a
# wedged box (Docker/journald/postgres all misbehave badly at 100% full).
#
# A root cron checks free space every 10 minutes and removes the ballast if it
# drops below 10%. This is a ONE-SHOT valve: once spent, recreate it with the
# same fallocate command below.

step "Emergency Disk Ballast"
if ! skip; then
  DOCKERS_DIR="$USER_HOME/dockers"
  BALLAST_FILE="$DOCKERS_DIR/emergency_empty_file.log"
  LOCAL_DIR="$USER_HOME/.local"
  LOCAL_BIN="$LOCAL_DIR/bin"
  CHECK_SCRIPT="$LOCAL_BIN/emergency-disk-check"
  CRON_FILE="/etc/cron.d/emergency-disk-check"

  # Name $LOCAL_DIR explicitly. `install -d` applies -o/-g ONLY to the directories given as
  # operands — any parent it has to invent on the way gets the caller's ownership, which here
  # is root. This step is the first thing to touch ~/.local, so leaving it implicit created a
  # root:root ~/.local with a user-owned bin/ inside it, and locked the user out of their own
  # ~/.local: anything unpacking or writing directly there (a Go tarball, npm's --prefix,
  # GOPATH) fails with a bare "Permission denied" far from this line.
  install -d -o "$USERNAME" -g "$USERNAME" "$DOCKERS_DIR" "$LOCAL_DIR" "$LOCAL_BIN"

  # --- the ballast file ---
  if [[ -f "$BALLAST_FILE" ]]; then
    ok "Ballast already exists at $BALLAST_FILE ($(du -h "$BALLAST_FILE" | cut -f1))"
    SUMMARY+=("Disk ballast: already present")
  else
    info "Allocating ballast file (10% of available disk)..."
    AVAIL_KB=$(df -Pk "$DOCKERS_DIR" | awk 'NR==2 {print $4}')
    BALLAST_KB=$(( AVAIL_KB / 10 ))

    # fallocate reserves real blocks (a sparse/truncate file would reserve
    # nothing and free nothing when deleted). dd is the fallback for
    # filesystems that don't support fallocate.
    if ! fallocate -l "${BALLAST_KB}K" "$BALLAST_FILE" 2>/dev/null; then
      warn "fallocate unsupported here — falling back to dd (slower)"
      dd if=/dev/zero of="$BALLAST_FILE" bs=1K count="$BALLAST_KB" status=none
    fi

    chown "$USERNAME:$USERNAME" "$BALLAST_FILE"
    chmod 644 "$BALLAST_FILE"
    ok "Ballast created: $BALLAST_FILE ($(du -h "$BALLAST_FILE" | cut -f1))"
    SUMMARY+=("Disk ballast: $(du -h "$BALLAST_FILE" | cut -f1) at $BALLAST_FILE")
  fi

  # --- the checker script ---
  info "Installing disk check script..."
  cat > "$CHECK_SCRIPT" <<'CHECKEOF'
#!/usr/bin/env bash
#
# Emergency disk ballast checker.
# Deletes the pre-allocated ballast file when free space falls below the
# threshold, buying headroom to log in and clean up. Installed by setup-ubuntu.sh.
#
# Run manually with --status to see current usage without changing anything.

set -euo pipefail

BALLAST="__BALLAST_PATH__"
THRESHOLD=10          # act when free space drops below this percentage
TAG="emergency-disk"

MOUNT_DIR="$(dirname "$BALLAST")"
USE_PCT="$(df -P "$MOUNT_DIR" | awk 'NR==2 {gsub(/%/,"",$5); print $5}')"
FREE_PCT=$(( 100 - USE_PCT ))

if [[ "${1:-}" == "--status" ]]; then
  echo "Mount:     $(df -P "$MOUNT_DIR" | awk 'NR==2 {print $6}')"
  echo "Free:      ${FREE_PCT}% (threshold: ${THRESHOLD}%)"
  if [[ -f "$BALLAST" ]]; then
    echo "Ballast:   present, $(du -h "$BALLAST" | cut -f1) — $BALLAST"
  else
    echo "Ballast:   ABSENT (already spent) — $BALLAST"
  fi
  exit 0
fi

(( FREE_PCT < THRESHOLD )) || exit 0

if [[ -f "$BALLAST" ]]; then
  FREED="$(du -h "$BALLAST" | cut -f1)"
  rm -f "$BALLAST"
  logger -t "$TAG" -p user.crit \
    "Free space ${FREE_PCT}% below ${THRESHOLD}% — deleted ballast, reclaimed ${FREED}. CLEAN UP NOW: this valve is spent."
else
  logger -t "$TAG" -p user.crit \
    "Free space ${FREE_PCT}% below ${THRESHOLD}% — ballast already spent, no headroom left to reclaim."
fi
CHECKEOF

  sed -i "s|__BALLAST_PATH__|$BALLAST_FILE|" "$CHECK_SCRIPT"
  chown "$USERNAME:$USERNAME" "$CHECK_SCRIPT"
  chmod 755 "$CHECK_SCRIPT"
  ok "Check script: $CHECK_SCRIPT"

  # --- the cron entry (root, so it can always delete) ---
  cat > "$CRON_FILE" <<EOF
# Emergency disk ballast — delete the ballast file if free space drops below 10%
# Installed by setup-ubuntu.sh. Check status: $CHECK_SCRIPT --status
PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
*/10 * * * * root $CHECK_SCRIPT
EOF
  chmod 644 "$CRON_FILE"
  ok "Cron installed: every 10 minutes ($CRON_FILE)"

  SUMMARY+=("Disk watchdog: $CHECK_SCRIPT via cron every 10m")
  step_ok
fi

# =============================================================================
# 23. Summary
# =============================================================================

echo ""
echo ""
if [[ ${#ERRORS[@]} -gt 0 ]]; then
  echo -e "${YELLOW}╔══════════════════════════════════════════════════╗${NC}"
  echo -e "${YELLOW}║  Setup Complete (with warnings)                  ║${NC}"
  echo -e "${YELLOW}╚══════════════════════════════════════════════════╝${NC}"
else
  echo -e "${GREEN}╔══════════════════════════════════════════════════╗${NC}"
  echo -e "${GREEN}║  Setup Complete                                  ║${NC}"
  echo -e "${GREEN}╚══════════════════════════════════════════════════╝${NC}"
fi

echo ""
echo -e "${BOLD}  What was done:${NC}"
for item in "${SUMMARY[@]}"; do
  echo -e "    ${GREEN}+${NC} $item"
done

if [[ ${#ERRORS[@]} -gt 0 ]]; then
  echo ""
  echo -e "${BOLD}  Non-critical issues:${NC}"
  for err in "${ERRORS[@]}"; do
    echo -e "    ${YELLOW}!${NC} $err"
  done
fi

echo ""
echo -e "${BOLD}  Server info:${NC}"
echo "    User:       $USERNAME"
echo "    Home:       $USER_HOME"
[[ -n "${TS_IP:-}" && "$TS_IP" != "unknown" ]] && echo "    Tailscale:  $TS_IP"
echo "    Shell:      zsh (Oh My Zsh + Starship)"

echo ""
echo -e "${BOLD}  Next steps:${NC}"
echo "    1. Reconnect as the new user:"
echo -e "       ${CYAN}ssh $USERNAME@<server-ip>${NC}"
echo "    2. Docker group requires a new login session"
echo ""

# Clean up progress file on success
rm -f "$PROGRESS_FILE"