#!/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"
# shellcheck source=lib/tools.sh
source "$SCRIPT_DIR/lib/tools.sh"
# shellcheck source=lib/system.sh
source "$SCRIPT_DIR/lib/system.sh"
# shellcheck source=lib/disk.sh
source "$SCRIPT_DIR/lib/disk.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"

# Always, and outside any step: everything below reads this index — core utils,
# the fastfetch PPA, the Docker repo — and `step` skips a step whose name is
# already in the progress file. With the refresh inside one of those, a resumed
# run installed against whatever the index happened to say hours or days ago.
echo ""
info "Refreshing the package index..."
pkg_refresh >/dev/null

# =============================================================================
# 2. Disk space
# =============================================================================
#
# First of the sections that change anything, because everything after it sizes
# itself from free disk — the swapfile and the ballast both.

step "Disk space"
if ! skip; then
  ROOT_DEV="$(root_device)"
  ROOT_DISK="$(parent_disk "$ROOT_DEV")"
  ROOT_PART="$(backing_partition "$ROOT_DEV" || true)"

  DISK_B="$(dev_bytes "$ROOT_DISK")"
  CONT_B="$(dev_bytes "$ROOT_DEV")"
  FS_B="$(fs_bytes "$ROOT_DEV")"
  VG_FREE_B="$(vg_free_bytes)"

  # A filesystem is always a little smaller than the thing holding it —
  # metadata, journal, reserved blocks. Only a real gap is worth reporting.
  SLACK=$((1024 * 1024 * 1024))

  echo ""
  info "Disk space — whether the root filesystem is actually using the whole drive"
  echo "     Ubuntu's installer, left on its defaults, gives the root volume a"
  echo "     fixed size and leaves the rest of the drive unallocated. On a 2TB"
  echo "     disk that is a 100G root and no sign anything is wrong: lsblk shows"
  echo "     the whole drive, df shows 100G, and the two are never seen together"
  echo "     until the day it fills up. Growing a virtual disk at the provider"
  echo "     leaves the same shape."
  echo ""
  echo "     drive:        $(human_bytes "$DISK_B")   ${ROOT_DISK}"
  echo "     volume:       $(human_bytes "$CONT_B")   ${ROOT_DEV}"
  echo "     filesystem:   $(human_bytes "$FS_B")   $(root_fstype), mounted at /"
  ((VG_FREE_B > SLACK)) && echo "     unused in LVM: $(human_bytes "$VG_FREE_B")"

  # growpart is the authority on whether a partition can move, but installing a
  # package just to ask is too eager — the arithmetic decides whether it is even
  # worth looking.
  MIGHT_GROW=false
  ((DISK_B - CONT_B > SLACK)) && MIGHT_GROW=true
  [[ -n "$ROOT_PART" ]] && $MIGHT_GROW && ensure_growpart

  if ((VG_FREE_B > SLACK)); then
    echo ""
    echo "     $(human_bytes "$VG_FREE_B") is sitting unallocated in the volume group"
    echo "     the fix is lvextend, then growing the filesystem into it — both online"
    if confirm "Use it?"; then
      grow_lv && grow_fs
      ok "root filesystem is now $(human_bytes "$(fs_bytes "$ROOT_DEV")")"
      SUMMARY+=("Disk: reclaimed $(human_bytes "$VG_FREE_B") from the volume group")
    else
      warn "skipped by request"
      SUMMARY+=("Disk: SKIPPED — $(human_bytes "$VG_FREE_B") left unallocated")
    fi

  elif [[ -n "$ROOT_PART" ]] && partition_can_grow "$ROOT_PART"; then
    echo ""
    echo "     the partition stops short of the end of the drive"
    echo "     this is the one step that edits the partition table — it only ever"
    echo "     moves the end of ${ROOT_PART} outwards, and never touches another one"
    if confirm "Extend it?"; then
      grow_partition "$ROOT_PART"
      if root_is_lvm; then
        grow_pv "$ROOT_PART"
        grow_lv
      fi
      grow_fs
      ok "root filesystem is now $(human_bytes "$(fs_bytes "$ROOT_DEV")")"
      SUMMARY+=("Disk: partition extended, filesystem now $(human_bytes "$(fs_bytes "$ROOT_DEV")")")
    else
      warn "skipped by request"
      SUMMARY+=("Disk: SKIPPED — partition left short of the drive")
    fi

  elif ((FS_B > 0 && CONT_B - FS_B > SLACK)); then
    echo ""
    echo "     the filesystem is smaller than the volume holding it"
    if confirm "Grow it?"; then
      grow_fs
      ok "root filesystem is now $(human_bytes "$(fs_bytes "$ROOT_DEV")")"
      SUMMARY+=("Disk: filesystem grown to $(human_bytes "$(fs_bytes "$ROOT_DEV")")")
    else
      warn "skipped by request"
      SUMMARY+=("Disk: SKIPPED — filesystem left short of its volume")
    fi

  else
    echo ""
    echo "     the whole drive is in use, nothing to reclaim"
    SUMMARY+=("Disk: already using the whole drive")
  fi
  step_ok
fi

# =============================================================================
# 3. System update
# =============================================================================
#
# Its own section because it is the only thing in the script that moves versions
# of software already on the machine. Everything else only ever adds what is
# absent, so this is the one that deserves to be refused on its own.
#
# The index refresh is NOT here — it runs in pre-flight, unconditionally, because
# every later step reads it and this one can be skipped.

step "System update"
if ! skip; then
  mapfile -t UPGRADABLE < <(pkg_upgradable)

  echo ""
  info "System update — upgrades installed packages to their latest versions"

  if ((${#UPGRADABLE[@]} == 0)); then
    echo "     to upgrade:   nothing, everything is current"
    SUMMARY+=("System update: already up to date")
  else
    echo "     to upgrade:   ${#UPGRADABLE[@]} package(s)"
    # Capped, because a box that has not been touched in months lists hundreds
    # and a wall of names is no more informative than a count.
    printf '       %s\n' "${UPGRADABLE[@]:0:25}"
    ((${#UPGRADABLE[@]} > 25)) && echo "       … and $((${#UPGRADABLE[@]} - 25)) more"

    if confirm "Proceed?"; then
      pkg_upgrade_all
      ok "System upgraded"
      SUMMARY+=("System upgraded: ${#UPGRADABLE[@]} package(s)")
    else
      warn "skipped by request"
      SUMMARY+=("System update: SKIPPED by request — ${#UPGRADABLE[@]} package(s) left as they are")
    fi
  fi
  step_ok
fi

# =============================================================================
# 4. Core utils
# =============================================================================
#
# What the distribution provides: the six this script would break without, and
# the command-line tools that make a machine worth sitting at.

step "Core utils"
if ! skip; then
  # shellcheck disable=SC2046  # word splitting is how the list is passed
  pkg_install "Core utils" $(pkgs_core)
  summarise_last "Core utils"
  step_ok
fi

# =============================================================================
# 5. Command-line tools
# =============================================================================
#
# A different thing from core utils, and kept apart from them: upstream binaries
# fetched from upstream, on their own release cadence, none of which the
# distribution ships. Lumping them in made a run look like it was installing
# system packages and then start pulling tarballs unannounced.

step "Command-line tools"
if ! skip; then
  # shellcheck disable=SC2046  # word splitting is how the list is passed
  tools_install "Command-line tools" $(tools_default)
  summarise_last "Command-line tools"
  step_ok
fi


# =============================================================================
# 6. Locale
# =============================================================================
#
# LOCALE in the environment overrides the default.

step "Locale"
if ! skip; then
  LOCALE="${LOCALE:-en_US.UTF-8}"
  CURRENT_LOCALE="$(locale_current)"

  echo ""
  info "Locale — the system language and character encoding"
  echo "     current:      ${CURRENT_LOCALE:-none configured}"
  echo "     to set:       ${LOCALE}"

  if [[ "$CURRENT_LOCALE" == "$LOCALE" ]] && locale_is_generated "$LOCALE"; then
    echo "     already set and generated, nothing to do"
    SUMMARY+=("Locale: already ${LOCALE}")
  else
    # Say which of the two is actually wrong, since they fail differently: a
    # missing LANG means the C locale, a missing generation means every login
    # prints a setlocale warning.
    [[ "$CURRENT_LOCALE" != "$LOCALE" ]] && echo "     LANG is not set to it"
    locale_is_generated "$LOCALE" || echo "     the locale has not been generated on this machine"

    if confirm "Proceed?"; then
      locale_set "$LOCALE"
      ok "Locale set to ${LOCALE}"
      SUMMARY+=("Locale: ${LOCALE}")
    else
      warn "skipped by request"
      SUMMARY+=("Locale: SKIPPED by request — left at ${CURRENT_LOCALE:-unset}")
    fi
  fi
  step_ok
fi

# =============================================================================
# 7. Timezone
# =============================================================================
#
# TIMEZONE in the environment answers the prompt ahead of time.

step "Timezone"
if ! skip; then
  CURRENT_TZ="$(timezone_current)"

  echo ""
  info "Timezone — what logs, timers and every printed date are relative to"
  echo "     current:      ${CURRENT_TZ:-unknown}"

  if [[ -z "${TIMEZONE:-}" ]]; then
    echo ""
    for i in "${!TZ_OPTIONS[@]}"; do
      printf '     [%d] %s\n' "$((i + 1))" "${TZ_OPTIONS[$i]}"
    done
    echo ""

    while [[ -z "${TIMEZONE:-}" ]]; do
      if ! read -rp "  Pick a number, or type a zone name — Enter keeps ${CURRENT_TZ:-the current one}: " TZ_CHOICE; then
        echo ""
        fail "No answer. Set TIMEZONE=<zone> to answer this ahead of time."
      fi

      if [[ -z "$TZ_CHOICE" ]]; then
        TIMEZONE="$CURRENT_TZ"
      elif [[ "$TZ_CHOICE" =~ ^[0-9]+$ ]]; then
        if ((TZ_CHOICE >= 1 && TZ_CHOICE <= ${#TZ_OPTIONS[@]})); then
          TIMEZONE="${TZ_OPTIONS[$((TZ_CHOICE - 1))]}"
        else
          warn "There is no option ${TZ_CHOICE}."
        fi
      else
        # Validated here rather than left to timedatectl, which fails on an
        # unknown name and would take the whole run down over a typo.
        if timezone_is_valid "$TZ_CHOICE"; then
          TIMEZONE="$TZ_CHOICE"
        else
          warn "Not a zone this machine knows: '${TZ_CHOICE}' — try e.g. Europe/Berlin"
        fi
      fi
    done
  elif ! timezone_is_valid "$TIMEZONE"; then
    fail "TIMEZONE='${TIMEZONE}' is not a zone this machine knows."
  fi

  if [[ "$TIMEZONE" == "$CURRENT_TZ" ]]; then
    echo "     keeping ${CURRENT_TZ}, nothing to do"
    SUMMARY+=("Timezone: already ${CURRENT_TZ}")
  else
    echo "     to set:       ${TIMEZONE}"
    if confirm "Proceed?"; then
      timezone_set "$TIMEZONE"
      ok "Timezone set to ${TIMEZONE}"
      SUMMARY+=("Timezone: ${TIMEZONE}")
    else
      warn "skipped by request"
      SUMMARY+=("Timezone: SKIPPED by request — left at ${CURRENT_TZ:-unknown}")
    fi
  fi
  step_ok
fi

# =============================================================================
# 8. Swap
# =============================================================================
#
# Disk the kernel can park cold pages on when RAM fills, so a spike costs
# latency instead of a process. A `bun install` or a Docker build on a small
# machine is exactly the spike this is for.

step "Swap"
if ! skip; then
  ACTIVE_SWAP_GB="$(swap_active_gb)"
  WANT_SWAP_GB="$(swap_recommended_gb)"
  SWAPPINESS="$(swappiness_for_role)"
  CURRENT_SWAPPINESS="$(sysctl -n vm.swappiness 2>/dev/null || echo unknown)"

  echo ""
  info "Swap — overflow space so a memory spike costs speed rather than a process"
  echo "     RAM:          $(ram_gb)G"
  echo "     active swap:  ${ACTIVE_SWAP_GB}G"
  echo "     swappiness:   ${CURRENT_SWAPPINESS} -> ${SWAPPINESS} (${MACHINE_ROLE})"

  if [[ "$IS_WSL" == true ]]; then
    # WSL2 runs its own managed swap inside the VM; a swapfile here is wasted
    # disk and is not what the kernel would use anyway.
    echo "     WSL manages its own swap — leaving it alone"
    SUMMARY+=("Swap: left to WSL")
  elif ((ACTIVE_SWAP_GB > 0)); then
    echo "     already has ${ACTIVE_SWAP_GB}G of swap, leaving it alone"
    if [[ "$CURRENT_SWAPPINESS" != "$SWAPPINESS" ]] && confirm "Set swappiness to ${SWAPPINESS}?"; then
      swappiness_set "$SWAPPINESS"
      ok "swappiness set to ${SWAPPINESS}"
      SUMMARY+=("Swap: kept ${ACTIVE_SWAP_GB}G, swappiness ${SWAPPINESS}")
    else
      SUMMARY+=("Swap: kept ${ACTIVE_SWAP_GB}G")
    fi
  elif ((WANT_SWAP_GB == 0)); then
    # Capped to nothing by the disk check rather than by choice.
    warn "not enough free disk to add swap safely — $(disk_free_gb)G free"
    SUMMARY+=("Swap: none added, disk too full")
  else
    echo "     to create:    ${WANT_SWAP_GB}G at ${SWAPFILE} ($(disk_free_gb)G free now)"
    if confirm "Proceed?"; then
      swap_create "$WANT_SWAP_GB"
      swappiness_set "$SWAPPINESS"
      ok "${WANT_SWAP_GB}G swap active, swappiness ${SWAPPINESS}"
      SUMMARY+=("Swap: ${WANT_SWAP_GB}G created, swappiness ${SWAPPINESS}")
    else
      warn "skipped by request"
      SUMMARY+=("Swap: SKIPPED by request")
    fi
  fi
  step_ok
fi

# =============================================================================
# 9. Emergency disk ballast
# =============================================================================
#
# Always offered, whatever the role — the role only decides which way the
# recommendation points.

step "Emergency disk ballast"
if ! skip; then
  echo ""
  info "Emergency disk ballast — a reserve you can burn when the disk fills up"
  echo "     A junk file holding no data, sized at 10% of free disk. A root cron"
  echo "     checks every 10 minutes and deletes it if free space drops below"
  echo "     ${BALLAST_THRESHOLD}%, so you get room to log in and clean up instead of meeting a"
  echo "     wedged machine — Docker, journald and postgres all misbehave badly"
  echo "     at 100% full, and not all of them recover on their own."
  echo "     It is a one-shot valve: once spent, run this again to recreate it."
  echo ""
  if is_server; then
    echo "     recommended for ${MACHINE_ROLE} — a full disk on an unattended box is the bad case"
  else
    echo "     less useful on ${MACHINE_ROLE} — you are sitting at this machine and will notice"
  fi

  if ballast_exists; then
    echo "     already present: $(ballast_size_human) at ${BALLAST_FILE}"
    SUMMARY+=("Disk ballast: already present ($(ballast_size_human))")
  else
    BALLAST_MB="$(ballast_size_mb)"
    echo "     to create:    ${BALLAST_MB}M at ${BALLAST_FILE}"
    if confirm "Create it?"; then
      ballast_create "$BALLAST_MB"
      ballast_install_checker
      ok "ballast $(ballast_size_human), checker at ${BALLAST_CHECKER}"
      SUMMARY+=("Disk ballast: $(ballast_size_human), checked every 10 min")
    else
      warn "skipped by request"
      SUMMARY+=("Disk ballast: SKIPPED by request")
    fi
  fi
  step_ok
fi

# =============================================================================
# 10. earlyoom
# =============================================================================

step "earlyoom"
if ! skip; then
  echo ""
  info "earlyoom — keeps the machine reachable when it runs out of memory"
  echo "     The kernel's own OOM killer waits until an allocation actually"
  echo "     fails, and by then the machine has usually spent minutes thrashing:"
  echo "     unresponsive, ssh refusing to connect, nothing to do but reset it."
  echo "     earlyoom watches free memory and kills the biggest consumer while"
  echo "     there is still enough left to stay logged in."
  echo ""
  echo "     This is what happens after swap runs out, so the two go together."

  if earlyoom_is_active; then
    echo "     already installed and running"
    SUMMARY+=("earlyoom: already running")
  elif confirm "Install it?"; then
    earlyoom_install
    if earlyoom_is_active; then
      ok "earlyoom running"
      SUMMARY+=("earlyoom: installed and running")
    else
      warn "earlyoom installed but not running — check: systemctl status earlyoom"
      SUMMARY+=("earlyoom: installed, not running")
    fi
  else
    warn "skipped by request"
    SUMMARY+=("earlyoom: SKIPPED by request")
  fi
  step_ok
fi

# =============================================================================
# 11. inotify watch limit
# =============================================================================

step "inotify watch limit"
if ! skip; then
  CURRENT_WATCHES="$(inotify_current_watches)"

  echo ""
  info "inotify watch limit — how many files can be watched for changes at once"
  echo "     A single file watcher walking a project with node_modules in it can"
  echo "     exhaust the stock limit on its own, and every watcher on the machine"
  echo "     draws from the same pool. The failure is silent: nothing errors, the"
  echo "     watcher just stops noticing changes. Hot reload goes quiet, a build"
  echo "     stops rebuilding, and the reason is never on screen."
  echo ""
  echo "     current:      ${CURRENT_WATCHES}"
  echo "     to set:       ${INOTIFY_WATCHES}"
  if is_role dev; then
    echo "     recommended on dev — editors, bun --watch and vite are all watchers"
  else
    echo "     less pressing on ${MACHINE_ROLE}, but anything running bun --watch or"
    echo "     serving a file browser is a watcher too"
  fi

  if ((CURRENT_WATCHES >= INOTIFY_WATCHES)); then
    echo "     already at or above that, nothing to do"
    SUMMARY+=("inotify watches: already ${CURRENT_WATCHES}")
  elif confirm "Raise it?"; then
    inotify_raise
    ok "inotify watches raised to $(inotify_current_watches)"
    SUMMARY+=("inotify watches: raised to ${INOTIFY_WATCHES}")
  else
    warn "skipped by request"
    SUMMARY+=("inotify watches: SKIPPED by request — left at ${CURRENT_WATCHES}")
  fi
  step_ok
fi

# =============================================================================
# NOT PORTED YET
# =============================================================================
#
# Sections still to move across from scripts/setup-old/setup-ubuntu.sh, in order:
#
#   auto-suspend · boot-hang fix · user creation ·
#   ssh keys · ssh hardening · dns · static ip · fail2ban · unattended-upgrades ·
#   git config · docker · zsh + prompt · tailscale · neovim · js runtimes ·
#   dev tools · ufw · zshrc
#
# Each arrives as its own commit. Delete this block when the list is empty.


# =============================================================================
# 12. 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}  Machine:${NC}"
echo "    System:     $OS_NAME ($ARCH)"
echo "    Role:       $MACHINE_ROLE"
echo "    User:       $USERNAME"
echo "    Home:       $USER_HOME"
[[ -n "${TS_IP:-}" && "$TS_IP" != "unknown" ]] && echo "    Tailscale:  $TS_IP"
echo ""

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