#!/usr/bin/env bash

# =============================================================================
#  Atlas Desktop Installer
#  for Arch Linux
# =============================================================================
# By Gu://em_
# Co-Authored with Claude Sonnet 4.6

set -euo pipefail

# ── Constants ─────────────────────────────────────────────────────────────────
REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
CONFIG_DIR="$REPO_DIR/config"
PACKAGES_DIR="$REPO_DIR/packages"
SCRIPTS_DIR="$REPO_DIR/scripts"
DOTFILES_DIR="$HOME/.atlas-dotfiles"
STATE_DIR="$DOTFILES_DIR/.state"
PKG_STATE_DIR="$STATE_DIR/packages"
WAYBAR_CFG_DIR="$HOME/.config/waybar"

# ── Colours & styles ──────────────────────────────────────────────────────────
ESC=$'\e'
BOLD="${ESC}[1m"      DIM="${ESC}[2m"       RESET="${ESC}[0m"
BLACK="${ESC}[30m"    RED="${ESC}[31m"       GREEN="${ESC}[32m"
YELLOW="${ESC}[33m"   BLUE="${ESC}[34m"      MAGENTA="${ESC}[35m"
CYAN="${ESC}[36m"     WHITE="${ESC}[37m"
BG_BLACK="${ESC}[40m" BG_BLUE="${ESC}[44m"
HIDE_CURSOR="${ESC}[?25l"  SHOW_CURSOR="${ESC}[?25h"
CLEAR_LINE="${ESC}[2K\r"

# Ensure cursor is restored on exit
trap 'printf "%s" "$SHOW_CURSOR"; tput cnorm 2>/dev/null || true' EXIT
trap 'printf "%s" "$SHOW_CURSOR"; tput cnorm 2>/dev/null || true; exit 1' INT TERM

# ── Logging helpers ───────────────────────────────────────────────────────────
PHASE_NUM=0

ok()      { printf "  ${GREEN}${BOLD}✔${RESET}  %s\n" "$*"; }
info()    { printf "  ${CYAN}·${RESET}  %s\n" "$*"; }
warn()    { printf "  ${YELLOW}⚠${RESET}  %s\n" "$*"; }
err()     { printf "  ${RED}${BOLD}✘${RESET}  %s\n" "$*" >&2; }
die()     { err "$*"; exit 1; }
blank()   { echo; }

indent()  { sed 's/^/     /'; }

phase() {
  (( PHASE_NUM++ )) || true
  blank
  printf "${BOLD}${BLUE}┌─ Phase %d ─ %s${RESET}\n" "$PHASE_NUM" "$*"
}

step() {
  printf "${BLUE}│${RESET}  ${BOLD}%s${RESET}\n" "$*"
}

phase_done() {
  printf "${BLUE}└─${GREEN} done${RESET}\n"
}

# ── Spinner ───────────────────────────────────────────────────────────────────
_SPINNER_PID=""
_SPINNER_FRAMES=('⠋' '⠙' '⠹' '⠸' '⠼' '⠴' '⠦' '⠧' '⠇' '⠏')

spinner_start() {
  local msg="$1"
  printf "%s" "$HIDE_CURSOR"
  (
    local i=0
    while true; do
      printf "${CLEAR_LINE}  ${CYAN}${_SPINNER_FRAMES[$((i % ${#_SPINNER_FRAMES[@]}))]}${RESET}  %s" "$msg"
      sleep 0.08
      (( i++ )) || true
    done
  ) &
  _SPINNER_PID=$!
}

spinner_stop() {
  local status="${1:-ok}"   # ok | fail
  if [[ -n $_SPINNER_PID ]]; then
    kill "$_SPINNER_PID" 2>/dev/null; wait "$_SPINNER_PID" 2>/dev/null || true
    _SPINNER_PID=""
  fi
  printf "%s" "$SHOW_CURSOR"
  printf "%s" "$CLEAR_LINE"
  if [[ $status == ok ]]; then
    ok "$2"
  else
    err "$2"
  fi
}

run_quiet() {
  # run_quiet "label" cmd [args...]
  local label="$1"; shift
  spinner_start "$label…"
  local tmp; tmp=$(mktemp)
  if "$@" >"$tmp" 2>&1; then
    spinner_stop ok "$label"
    rm -f "$tmp"
  else
    local rc=$?
    spinner_stop fail "$label"
    warn "Output:"
    cat "$tmp" | indent
    rm -f "$tmp"
    return $rc
  fi
}

# ── UI primitives ─────────────────────────────────────────────────────────────
header() {
  clear
  printf "${BOLD}${BLUE}"
  cat << 'EOF'
  ╔══════════════════════════════════════════════════════════════════╗
  ║          _   _             _____            _    _               ║
  ║     /\  | | | |           |  __ \          | |  | |              ║
  ║    /  \ | |_| | __ _ ___  | |  | | ___  ___| | _| |_ ___  _ __   ║
  ║   / /\ \| __| |/ _` / __| | |  | |/ _ \/ __| |/ / __/ _ \| '_ \  ║
  ║  / ____ \ |_| | (_| \__ \ | |__| |  __/\__ \   <| || (_) | |_) | ║
  ║ /_/    \_\__|_|\__,_|___/ |_____/ \___||___/_|\_\\__\___/| .__/  ║
  ║                                                          | |     ║
  ║                                                          |_|     ║
  ╠══════════════════════════════════════════════════════════════════╣
  ║                     Installer  ·  Arch Linux                     ║
  ╚══════════════════════════════════════════════════════════════════╝
EOF
  printf "${RESET}\n"
}

divider() {
  printf "  ${DIM}%s${RESET}\n" "──────────────────────────────────────────────────────"
}

confirm() {
  local msg="$1" default="${2:-y}"
  local prompt; [[ $default == y ]] && prompt="${GREEN}Y${RESET}/n" || prompt="y/${RED}N${RESET}"
  local ans
  read -rp "  ${YELLOW}?${RESET}  $msg [${prompt}] " ans
  ans="${ans:-$default}"
  [[ ${ans,,} == y ]]
}

press_enter() {
  read -rp "  ${DIM}Press Enter to continue…${RESET}" _
  blank
}

# ── State helpers ─────────────────────────────────────────────────────────────
state_get() { grep -m1 "^${1}=" "$STATE_DIR/install.state" 2>/dev/null | cut -d= -f2- || true; }
state_set() {
  local key="$1" val="$2"
  local f="$STATE_DIR/install.state"
  if grep -q "^${key}=" "$f" 2>/dev/null; then
    sed -i "s|^${key}=.*|${key}=${val}|" "$f"
  else
    echo "${key}=${val}" >> "$f"
  fi
}

# ── Package helpers ───────────────────────────────────────────────────────────
read_pkg_file() {
  local file="$1"
  [[ -f $file ]] || return
  grep -v '^\s*#' "$file" | grep -v '^\s*$' | sort -u
}

save_pkg_list() {
  local group="$1"; shift
  mkdir -p "$PKG_STATE_DIR"
  printf '%s\n' "$@" > "$PKG_STATE_DIR/${group}"
}

pacman_install() {
  sudo pacman -S --needed --noconfirm "$@"
}

paru_install() {
  paru -S --needed --noconfirm "$@"
}

pkg_installed() {
  pacman -Q "$1" &>/dev/null
}

# ── Hardware detection ────────────────────────────────────────────────────────
detect_hardware() {
  HW_GPU="unknown"
  HW_CPU="unknown"

  if grep -qi intel /proc/cpuinfo 2>/dev/null; then HW_CPU="intel"; fi
  if grep -qi amd   /proc/cpuinfo 2>/dev/null; then HW_CPU="amd"; fi

  if lspci 2>/dev/null | grep -qi "nvidia";                        then HW_GPU="nvidia"
  elif lspci 2>/dev/null | grep -qi "amd.*radeon\|advanced micro"; then HW_GPU="amd"
  elif lspci 2>/dev/null | grep -qi "intel.*graphics";             then HW_GPU="intel"
  fi
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 0 — Preflight
# ─────────────────────────────────────────────────────────────────────────────
preflight() {
  phase "Preflight checks"

  # Arch Linux check
  step "Checking system"
  if [[ ! -f /etc/arch-release ]]; then
    err "This installer only supports Arch Linux."
    exit 1
  fi
  ok "Arch Linux detected"

  # Not root
  if [[ $EUID -eq 0 ]]; then
    die "Do not run this script as root. It will use sudo when needed."
  fi
  ok "Running as user: ${BOLD}${USER}${RESET}"

  # sudo available
  if ! sudo -v 2>/dev/null; then
    die "sudo is required. Make sure $USER is in the sudoers group."
  fi
  ok "sudo access confirmed"

  # git
  command -v git &>/dev/null || die "git is not installed."
  ok "git found"

  # Config dir
  if [[ ! -d $CONFIG_DIR ]]; then
    die "config/ directory not found in $REPO_DIR"
  fi
  ok "Repo looks good: ${DIM}$REPO_DIR${RESET}"

  detect_hardware
  info "CPU: ${BOLD}${HW_CPU^^}${RESET}   GPU: ${BOLD}${HW_GPU^^}${RESET}"

  phase_done
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 1 — Welcome & choices
# ─────────────────────────────────────────────────────────────────────────────
welcome() {
  header
  cat << EOF
  Welcome! This script will set up the Atlas Desktop on your machine.
  It will:

  ${CYAN}·${RESET} Install required packages (pacman + paru for AUR)
  ${CYAN}·${RESET} Link your dotfiles into ${BOLD}~/${RESET} using GNU Stow
  ${CYAN}·${RESET} Enable necessary system services
  ${CYAN}·${RESET} Configure the updater so you stay in sync

  Nothing will be modified until you confirm below.

EOF
  divider
  blank
  confirm "Ready to begin?" y || { echo; info "Goodbye!"; exit 0; }
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 2 — Waybar layout picker
# ─────────────────────────────────────────────────────────────────────────────
choose_waybar() {
  phase "Waybar layout"
  blank

  # Vertical preview
  printf "  ${BOLD}${CYAN}1) Vertical${RESET}  ${DIM}(sidebar, left or right)${RESET}\n\n"
  printf "${DIM}"
  cat << 'EOF'
     ┌────┐
     │ 1  │  ← workspace indicators
     │ 2  │
     │    │
     │    │
     │ 󰕾  │  ← volume
     │ 󰖩  │  ← network
     │ 12 │
     │ 34 │  ← clock
     │   │  ← power
     └────┘
EOF
  printf "${RESET}\n"

  # Horizontal preview
  printf "  ${BOLD}${CYAN}2) Horizontal${RESET}  ${DIM}(bottom bar)${RESET}\n\n"
  printf "${DIM}"
  cat << 'EOF'
     ┌─────────────────────────────────────────────┐
     │ ● ● ○ ○             12:34          󰕾  󰖩    │
     └─────────────────────────────────────────────┘
EOF
  printf "${RESET}\n"

  divider
  local choice
  while true; do
    read -rp "  ${YELLOW}?${RESET}  Choose a layout ${BOLD}[1/2]${RESET}: " choice
    case "$choice" in
      1) WAYBAR_LAYOUT="vertical";   ok "Vertical layout selected";   break ;;
      2) WAYBAR_LAYOUT="horizontal"; ok "Horizontal layout selected"; break ;;
      *) warn "Please enter 1 or 2." ;;
    esac
  done
  state_set "waybar_layout" "$WAYBAR_LAYOUT"
  phase_done
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 3 — Security profile
# ─────────────────────────────────────────────────────────────────────────────
choose_security() {
  phase "Security & stability"
  blank
  cat << EOF
  The following extras are ${BOLD}recommended for most users${RESET} but are
  optional for those who manage their own setup:

  ${GREEN}·${RESET} ${BOLD}UFW${RESET}         — simple firewall (blocks unsolicited inbound traffic)
  ${GREEN}·${RESET} ${BOLD}AppArmor${RESET}    — mandatory access control (limits process capabilities)
  ${GREEN}·${RESET} ${BOLD}Zram${RESET}        — compressed swap-in-RAM (better performance, less disk wear)

  ${DIM}Recommended unless you already manage these yourself.${RESET}

EOF
  if confirm "Install security & stability extras?" y; then
    INSTALL_SECURITY=true
    ok "Security extras will be installed"
  else
    INSTALL_SECURITY=false
    info "Skipping — you can install them later"
  fi
  state_set "install_security" "$INSTALL_SECURITY"
  phase_done
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 4 — Optional package groups
# ─────────────────────────────────────────────────────────────────────────────
choose_optional_groups() {
  phase "Optional packages"
  blank

  local optional_file="$PACKAGES_DIR/optional.groups"
  if [[ ! -f $optional_file ]]; then
    info "No optional.groups file found — skipping"
    CHOSEN_OPTIONAL_GROUPS=()
    phase_done
    return
  fi

  mapfile -t ALL_OPTIONAL < <(grep -v '^\s*#' "$optional_file" | grep -v '^\s*$')

  if (( ${#ALL_OPTIONAL[@]} == 0 )); then
    info "No optional groups defined — skipping"
    CHOSEN_OPTIONAL_GROUPS=()
    phase_done
    return
  fi

  info "Select which optional groups to install:"
  blank

  CHOSEN_OPTIONAL_GROUPS=()
  for group in "${ALL_OPTIONAL[@]}"; do
    # Show a short description if <group>.desc file exists
    local desc_file="$PACKAGES_DIR/${group}.desc"
    local desc=""
    [[ -f $desc_file ]] && desc=" ${DIM}— $(cat "$desc_file")${RESET}"
    if confirm "  Install ${BOLD}${group}${RESET}${desc}?" n; then
      CHOSEN_OPTIONAL_GROUPS+=("$group")
    fi
  done

  if (( ${#CHOSEN_OPTIONAL_GROUPS[@]} )); then
    ok "Selected: ${CHOSEN_OPTIONAL_GROUPS[*]}"
  else
    info "No optional groups selected"
  fi

  # Save choices
  local joined; joined=$(IFS=','; echo "${CHOSEN_OPTIONAL_GROUPS[*]}")
  state_set "optional_groups" "$joined"
  phase_done
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 5 — Summary & confirm
# ─────────────────────────────────────────────────────────────────────────────
confirm_summary() {
  phase "Summary"
  blank
  printf "  Here's what will happen:\n\n"

  printf "  ${BOLD}Dotfiles${RESET}\n"
  printf "  ${DIM}  %s  →  %s${RESET}\n" "$CONFIG_DIR" "$HOME"
  blank

  printf "  ${BOLD}Waybar layout:${RESET} %s\n" "${WAYBAR_LAYOUT}"
  blank

  printf "  ${BOLD}Hardware packages:${RESET}\n"
  printf "    CPU: %s  ·  GPU: %s\n" "${HW_CPU}" "${HW_GPU}"
  blank

  if [[ $INSTALL_SECURITY == true ]]; then
    printf "  ${GREEN}✔${RESET}  Security extras: UFW, AppArmor, Zram\n"
  else
    printf "  ${DIM}·  Security extras: skipped${RESET}\n"
  fi
  blank

  if (( ${#CHOSEN_OPTIONAL_GROUPS[@]} )); then
    printf "  ${BOLD}Optional groups:${RESET} %s\n" "${CHOSEN_OPTIONAL_GROUPS[*]}"
  else
    printf "  ${DIM}·  No optional groups${RESET}\n"
  fi
  blank

  printf "  ${BOLD}Services to enable:${RESET}\n"
  printf "    ly.service"
  [[ $INSTALL_SECURITY == true ]] && printf "  ·  apparmor.service  ·  ufw  ·  systemd-zram-setup@zram0"
  printf "\n"
  blank

  divider
  blank
  confirm "Proceed with installation?" y || { info "Aborted — nothing was changed."; exit 0; }
  blank
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 6 — System update
# ─────────────────────────────────────────────────────────────────────────────
system_update() {
  phase "System update"
  step "Syncing pacman databases and updating system"
  run_quiet "Updating system" sudo pacman -Syu --noconfirm
  phase_done
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 7 — Bootstrap paru
# ─────────────────────────────────────────────────────────────────────────────
bootstrap_paru() {
  phase "AUR helper (paru)"

  if pkg_installed paru; then
    ok "paru is already installed"
    phase_done
    return
  fi

  step "Installing paru from AUR"
  info "Requires: base-devel, git, cargo (rust)"

  run_quiet "Installing base-devel" sudo pacman -S --needed --noconfirm base-devel git

  # Install rust if needed (paru needs cargo)
  if ! command -v cargo &>/dev/null; then
    run_quiet "Installing rust" sudo pacman -S --needed --noconfirm rust
  fi

  local tmp_dir; tmp_dir=$(mktemp -d)
  spinner_start "Cloning paru…"
  git clone --depth=1 https://aur.archlinux.org/paru.git "$tmp_dir/paru" &>/dev/null
  spinner_stop ok "Cloned paru"

  step "Building paru (this may take a moment)"
  (
    cd "$tmp_dir/paru"
    run_quiet "Building paru" makepkg -si --noconfirm
  )
  rm -rf "$tmp_dir"

  if pkg_installed paru; then
    ok "paru installed successfully"
  else
    die "paru installation failed"
  fi
  phase_done
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 8 — Install packages
# ─────────────────────────────────────────────────────────────────────────────
install_packages() {
  phase "Installing packages"

  local pkg_file pkgs group

  # Helper: install a group and save state
  # Note: AUR packages are installed one-by-one so a single failure does not abort
  # the whole group. Official repo packages are bulk-installed with a per-
  # package fallback. Only successfully installed packages are saved to state.
  install_group() {
    local group="$1"
    local pkg_file="$PACKAGES_DIR/${group}.pkgs"
    [[ -f $pkg_file ]] || { info "No file for group '${group}' — skipping"; return; }

    mapfile -t pkgs < <(read_pkg_file "$pkg_file")
    (( ${#pkgs[@]} == 0 )) && { info "${group}: empty — skipping"; return; }

    step "Group: ${group} (${#pkgs[@]} packages)"

    local succeeded=() failed=()

    if [[ $group == aur ]]; then
      for pkg in "${pkgs[@]}"; do
        local tmp; tmp=$(mktemp)
        if [[ $IS_TTY == true ]]; then
          spinner_start "  [aur] $pkg"
          if paru -S --needed --noconfirm "$pkg" >"$tmp" 2>&1; then
            spinner_stop ok "[aur] $pkg"; succeeded+=("$pkg")
          else
            spinner_stop fail "[aur] $pkg"
            warn "paru output:"; cat "$tmp" | indent; failed+=("$pkg")
          fi
        else
          printf "  ...  [aur] %s\n" "$pkg"
          if paru -S --needed --noconfirm "$pkg" >"$tmp" 2>&1; then
            ok "[aur] $pkg"; succeeded+=("$pkg")
          else
            err "[aur] $pkg"
            warn "paru output:"; cat "$tmp" | indent; failed+=("$pkg")
          fi
        fi
        rm -f "$tmp"
      done
    else
      run_quiet "  Installing ${group}" pacman_install "${pkgs[@]}"
      local tmp; tmp=$(mktemp)
      if pacman_install "${pkgs[@]}" >"$tmp" 2>&1; then
        succeeded=("${pkgs[@]}"); rm -f "$tmp"
      else
        rm -f "$tmp"
        warn "Bulk install failed for group '${group}' — retrying per package"
        for pkg in "${pkgs[@]}"; do
          local tmp2; tmp2=$(mktemp)
          if [[ $IS_TTY == true ]]; then
            spinner_start "  [${group}] $pkg"
            if pacman_install "$pkg" >"$tmp2" 2>&1; then
              spinner_stop ok "[${group}] $pkg"; succeeded+=("$pkg")
            else
              spinner_stop fail "[${group}] $pkg"
              warn "pacman output:"; cat "$tmp2" | indent; failed+=("$pkg")
            fi
          else
            printf "  ...  [%s] %s\n" "$group" "$pkg"
            if pacman_install "$pkg" >"$tmp2" 2>&1; then
              ok "[${group}] $pkg"; succeeded+=("$pkg")
            else
              err "[${group}] $pkg"
              warn "pacman output:"; cat "$tmp2" | indent; failed+=("$pkg")
            fi
          fi
          rm -f "$tmp2"
        done
      fi
    fi

    if (( ${#succeeded[@]} )); then
      save_pkg_list "$group" "${succeeded[@]}"
    fi

    if (( ${#failed[@]} )); then
      warn "Failed in group '${group}' (will retry on next update):"
      for pkg in "${failed[@]}"; do
        printf "       - %s\n" "$pkg"
      done
    fi
    return 0
  }

  # stow itself (needed for the next step)
  run_quiet "Ensuring stow is present" sudo pacman -S --needed --noconfirm stow

  # Build the set of groups that are handled specially so we can exclude them
  # when iterating all .pkgs files for automatic groups.
  local hw_groups=("${HW_CPU}-cpu" "${HW_GPU}-gpu")
  local optional_group_names=()
  local optional_file="$PACKAGES_DIR/optional.groups"
  if [[ -f $optional_file ]]; then
    mapfile -t optional_group_names < <(grep -v '^\s*#' "$optional_file" | grep -v '^\s*$')
  fi

  is_special_group() {
    local g="$1"
    # Returns 0 if the group is hw, optional, aur, or security (handled separately)
    [[ $g == aur || $g == security ]] && return 0
    for hw in "${hw_groups[@]}"; do [[ $g == "$hw" ]] && return 0; done
    for og in "${optional_group_names[@]}"; do [[ $g == "$og" ]] && return 0; done
    return 1
  }

  # 1. Core
  install_group "core"

  # 2. Hardware (auto-detected)
  install_group "${HW_CPU}-cpu"
  install_group "${HW_GPU}-gpu"

  # 3. All remaining automatic groups (every .pkgs file that is not special)
  while IFS= read -r pkg_file; do
    local group; group=$(basename "$pkg_file" .pkgs)
    is_special_group "$group" && continue
    [[ $group == core ]] && continue   # already done above
    install_group "$group"
  done < <(find "$PACKAGES_DIR" -maxdepth 1 -name "*.pkgs" | sort)


  # 4. Security extras
  if [[ $INSTALL_SECURITY == true ]]; then
    local sec_pkgs=(ufw apparmor zram-generator)
    step "Security extras"
    run_quiet "  Installing security packages" pacman_install "${sec_pkgs[@]}"
    save_pkg_list "security" "${sec_pkgs[@]}"
  fi

  # 5. Optional chosen groups
  for group in "${CHOSEN_OPTIONAL_GROUPS[@]}"; do
    install_group "$group"
  done

  # 6. AUR group
  install_group "aur"

  phase_done
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 9 — Clone into ~/.atlas-dotfiles and stow
# ─────────────────────────────────────────────────────────────────────────────
stow_dotfiles() {
  phase "Cloning and stowing dotfiles"

  # ~/.atlas-dotfiles is the permanent home of the repo. You can delete the original
  # clone after installation without breaking anything.
  local remote_url
  remote_url=$(git -C "$REPO_DIR" remote get-url origin 2>/dev/null || true)

  if [[ -z $remote_url ]]; then
    die "Could not determine remote URL from $REPO_DIR (is origin set?)"
  fi

  step "Setting up ~/.atlas-dotfiles as the permanent repo"
  info "Remote: $remote_url"

  if [[ -d $DOTFILES_DIR/.git ]]; then
    # Already a git repo — verify it points to the same remote
    local existing_remote
    existing_remote=$(git -C "$DOTFILES_DIR" remote get-url origin 2>/dev/null || true)
    if [[ $existing_remote == $remote_url ]]; then
      ok "~/.atlas-dotfiles already cloned from correct remote"
      run_quiet "Pulling latest" git -C "$DOTFILES_DIR" pull --ff-only origin main
    else
      warn "~/.atlas-dotfiles exists but points to a different remote:"
      warn "  existing: $existing_remote"
      warn "  expected: $remote_url"
      if confirm "Replace it with the correct repo?" y; then
        rm -rf "$DOTFILES_DIR"
        run_quiet "Cloning dotfiles" git clone "$remote_url" "$DOTFILES_DIR"
      else
        die "Aborted — ~/.atlas-dotfiles remote mismatch"
      fi
    fi
  elif [[ -e $DOTFILES_DIR ]]; then
    warn "~/.atlas-dotfiles exists but is not a git repo — moving to ~/.atlas-dotfiles.bak"
    mv "$DOTFILES_DIR" "${DOTFILES_DIR}.bak"
    run_quiet "Cloning dotfiles" git clone "$remote_url" "$DOTFILES_DIR"
  else
    run_quiet "Cloning dotfiles" git clone "$remote_url" "$DOTFILES_DIR"
  fi

  #TODO Checkout ?

  ok "~/.atlas-dotfiles → $remote_url"

  # Backup any real files that would conflict with stow links
  local stow_config_dir="$DOTFILES_DIR/config"
  local conflicts=()
  while IFS= read -r link; do
    local rel="${link#$stow_config_dir/}"
    local target="$HOME/$rel"
    if [[ -e $target && ! -L $target ]]; then
      conflicts+=("$target")
    fi
  done < <(find "$stow_config_dir" -type f 2>/dev/null)

  if (( ${#conflicts[@]} )); then
    warn "${#conflicts[@]} existing file(s) will be backed up (.bak):"
    for f in "${conflicts[@]}"; do
      info "  $f → ${f}.bak"
      mv "$f" "${f}.bak"
    done
  fi

  # Stow from ~/.atlas-dotfiles treating config/ as the stow package
  run_quiet "Stowing config" \
    stow --dir="$DOTFILES_DIR" --target="$HOME" --restow config

  phase_done
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 10 — Waybar config linking
# ─────────────────────────────────────────────────────────────────────────────
setup_waybar() {
  phase "Waybar configuration"

  local layout_dir="$WAYBAR_CFG_DIR/${WAYBAR_LAYOUT}"

  if [[ ! -d $layout_dir ]]; then
    warn "Waybar layout directory not found: $layout_dir"
    warn "Make sure config/.config/waybar/${WAYBAR_LAYOUT}/ exists in the repo"
    phase_done
    return
  fi

  step "Linking ${WAYBAR_LAYOUT} layout as active config"

  # Remove existing active config/style links (not the layout dirs themselves)
  rm -f "$WAYBAR_CFG_DIR/config" "$WAYBAR_CFG_DIR/style.css" \
        "$WAYBAR_CFG_DIR/config.jsonc"

  # Link config file (support config, config.json, config.jsonc)
  local cfg_file=""
  for name in config config.json config.jsonc; do
    [[ -f "$layout_dir/$name" ]] && { cfg_file="$name"; break; }
  done

  if [[ -n $cfg_file ]]; then
    ln -sf "$layout_dir/$cfg_file" "$WAYBAR_CFG_DIR/config"
    ok "config → ${WAYBAR_LAYOUT}/${cfg_file}"
  else
    warn "No config file found in $layout_dir"
  fi

  # Link style.css
  if [[ -f "$layout_dir/style.css" ]]; then
    ln -sf "$layout_dir/style.css" "$WAYBAR_CFG_DIR/style.css"
    ok "style.css → ${WAYBAR_LAYOUT}/style.css"
  else
    warn "No style.css found in $layout_dir"
  fi

  state_set "waybar_layout" "$WAYBAR_LAYOUT"
  phase_done
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 11 — Enable services
# ─────────────────────────────────────────────────────────────────────────────
enable_services() {
  phase "Enabling services"

  enable_service() {
    local svc="$1"
    if systemctl is-enabled "$svc" &>/dev/null; then
      ok "$svc (already enabled)"
    else
      if sudo systemctl enable "$svc" &>/dev/null; then
        ok "$svc"
      else
        warn "Failed to enable $svc (service may not be installed)"
      fi
    fi
  }

  # Display manager
  step "Display manager"
  # Disable any currently-enabled DM first to avoid conflicts
  for dm in gdm sddm lightdm lxdm greetd; do
    systemctl is-enabled "${dm}.service" &>/dev/null && \
      sudo systemctl disable "${dm}.service" &>/dev/null || true
  done
  enable_service "ly.service"

  # Security services (if chosen)
  if [[ $INSTALL_SECURITY == true ]]; then
    blank
    step "Security & stability"
    enable_service "apparmor.service"
    enable_service "ufw.service"

    # UFW default policy
    spinner_start "Configuring UFW defaults"
    sudo ufw default deny incoming  &>/dev/null
    sudo ufw default allow outgoing &>/dev/null
    sudo ufw --force enable         &>/dev/null
    spinner_stop ok "UFW configured (deny incoming, allow outgoing)"

    # Zram — write generator config if not present
    local zram_conf="/etc/systemd/zram-generator.conf"
    if [[ ! -f $zram_conf ]]; then
      spinner_start "Writing zram config"
      printf '[zram0]\nzram-size = ram / 2\ncompression-algorithm = zstd\n' \
        | sudo tee "$zram_conf" >/dev/null
      spinner_stop ok "Zram configured (50% RAM, zstd)"
    else
      ok "Zram config already exists"
    fi
    enable_service "systemd-zram-setup@zram0.service"
  fi

  # Pipewire (audio)
  blank
  step "Audio (Pipewire)"
  systemctl --user enable --now pipewire.service      &>/dev/null && ok "pipewire.service" || warn "pipewire.service not found"
  systemctl --user enable --now pipewire-pulse.service &>/dev/null && ok "pipewire-pulse.service" || warn "pipewire-pulse not found"
  systemctl --user enable --now wireplumber.service    &>/dev/null && ok "wireplumber.service" || warn "wireplumber not found"

  phase_done
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 12 — Install updater
# ─────────────────────────────────────────────────────────────────────────────
install_updater() {
  phase "Dotfiles updater"

  local updater_in_dotfiles="$DOTFILES_DIR/scripts/update.sh"
  local updater_link="$HOME/.local/bin/dotfiles-update"

  if [[ ! -f $updater_in_dotfiles ]]; then
    warn "update.sh not found in ~/.atlas-dotfiles/scripts/ — skipping"
    info "Expected: $updater_in_dotfiles"
    phase_done
    return
  fi

  chmod +x "$updater_in_dotfiles"
  mkdir -p "$HOME/.local/bin"

  # Symlink into PATH — stow may also handle this if scripts/ is a stow pkg,
  # but we do it explicitly here so it works regardless of stow layout.
  ln -sf "$updater_in_dotfiles" "$updater_link"
  ok "dotfiles-update → $updater_in_dotfiles"
  info "Run ${BOLD}dotfiles-update${RESET} any time to sync with upstream"

  # Initialise state inside the permanent repo
  local state_dir="$DOTFILES_DIR/.state"
  mkdir -p "$state_dir"
  # Point state file at the permanent repo, not the installer's location
  grep -q "^repo_dir=" "$state_dir/update.state" 2>/dev/null ||     echo "repo_dir=$DOTFILES_DIR" >> "$state_dir/update.state"
  # Stamp the current HEAD so the first `dotfiles-update` knows the baseline
  local head; head=$(git -C "$DOTFILES_DIR" rev-parse HEAD)
  if grep -q "^last_commit=" "$state_dir/update.state" 2>/dev/null; then
    sed -i "s|^last_commit=.*|last_commit=${head}|" "$state_dir/update.state"
  else
    echo "last_commit=$head" >> "$state_dir/update.state"
  fi

  phase_done
}

# ─────────────────────────────────────────────────────────────────────────────
#  STEP 13 — Final screen
# ─────────────────────────────────────────────────────────────────────────────
finish() {
  blank
  printf "${BOLD}${GREEN}"
  cat << 'EOF'
  ╔══════════════════════════════════════════════════════════╗
  ║                                                          ║
  ║     ✔  Installation complete!                            ║
  ║                                                          ║
  ╚══════════════════════════════════════════════════════════╝
EOF
  printf "${RESET}\n"

  cat << EOF
  ${BOLD}What to do next:${RESET}

  ${CYAN}·${RESET} ${BOLD}Reboot${RESET} to start Ly and load your Hyprland session
        ${DIM}sudo reboot${RESET}

  ${CYAN}·${RESET} ${BOLD}Update dotfiles${RESET} any time by running:
        ${DIM}atlas-update${RESET}

  ${CYAN}·${RESET} ${BOLD}Switch Waybar layout${RESET} (vertical ↔ horizontal):
        ${DIM}atlas-waybar [vertical|horizontal]${RESET}

  ${CYAN}·${RESET} Your chosen settings are saved in:
        ${DIM}$STATE_DIR/install.state${RESET}

EOF

  if [[ $INSTALL_SECURITY == true ]]; then
    printf "  ${GREEN}·${RESET}  ${BOLD}UFW${RESET} is active. To allow a port: ${DIM}sudo ufw allow <port>${RESET}\n"
    blank
  fi

  divider
  printf "  ${DIM}Atlas Desktop — %s${RESET}\n\n" "$(git -C "$REPO_DIR" rev-parse --short HEAD)"
}

# ─────────────────────────────────────────────────────────────────────────────
#  Waybar switcher (can be called directly: atlas-waybar vertical)
# ─────────────────────────────────────────────────────────────────────────────
waybar_switch() {
  local layout="${1:-}"
  if [[ -z $layout ]]; then
    local current; current=$(state_get "waybar_layout" || echo "unknown")
    echo "Current layout: $current"
    echo "Usage: atlas-waybar [vertical|horizontal]"
    return
  fi
  case "$layout" in
    vertical|horizontal) ;;
    *) die "Unknown layout '$layout'. Choose: vertical, horizontal" ;;
  esac
  WAYBAR_LAYOUT="$layout"
  setup_waybar
  ok "Waybar switched to ${layout}. Restart waybar to apply."
}

# ─────────────────────────────────────────────────────────────────────────────
#  Main
# ─────────────────────────────────────────────────────────────────────────────
main() {
  # Allow standalone waybar switching
  if [[ "${1:-}" == "waybar" ]]; then
    shift
    waybar_switch "$@"
    exit 0
  fi

  mkdir -p "$STATE_DIR" "$PKG_STATE_DIR"
  [[ -f "$STATE_DIR/install.state" ]] || touch "$STATE_DIR/install.state"

  # Init arrays
  CHOSEN_OPTIONAL_GROUPS=()
  INSTALL_SECURITY=false
  WAYBAR_LAYOUT="horizontal"
  HW_CPU="unknown"
  HW_GPU="unknown"

  welcome
  preflight
  choose_waybar
  choose_security
  choose_optional_groups
  confirm_summary

  # From here on we're actually making changes
  system_update
  bootstrap_paru
  install_packages
  stow_dotfiles
  setup_waybar
  enable_services
  install_updater

  finish
}

main "$@"