#!/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"

# 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. 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

# =============================================================================
# 3. 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

# =============================================================================
# 4. 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


# =============================================================================
# 5. 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

# =============================================================================
# 6. 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

# =============================================================================
# 7. 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

# =============================================================================
# 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 · disk ballast
#
# Each arrives as its own commit. Delete this block when the list is empty.


# =============================================================================
# 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}  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"