#!/usr/bin/env bash

# =============================================================================
#  Atlas Updater
#  Arch Linux · GNU Stow · pacman / paru
# =============================================================================
# By Gu://em_
# Co-Authored with Claude Sonnet 4.6

set -euo pipefail

# ── Paths ────────────────────────────────────────────────────────────────────
DOTFILES_DIR="${DOTFILES_DIR:-$HOME/.atlas-dotfiles}"
STATE_DIR="$DOTFILES_DIR/.state"
PKG_STATE_DIR="$STATE_DIR/packages"
STATE_FILE="$STATE_DIR/update.state"
PACKAGES_DIR="$DOTFILES_DIR/packages"
STOW_DIR="$DOTFILES_DIR"
STOW_PKG="config"
STOW_TARGET="$HOME"

# ── Colours ──────────────────────────────────────────────────────────────────
if [[ -t 1 ]]; then
  BOLD=$'\e[1m';    RESET=$'\e[0m'
  RED=$'\e[31m';    GREEN=$'\e[32m';  YELLOW=$'\e[33m'
  BLUE=$'\e[34m';   CYAN=$'\e[36m';   DIM=$'\e[2m'
else
  BOLD=''; RESET=''; RED=''; GREEN=''; YELLOW=''; BLUE=''; CYAN=''; DIM=''
fi

# ── Helpers ──────────────────────────────────────────────────────────────────
section()  { echo; echo "${BOLD}${BLUE}══ $* ══${RESET}"; }
ok()       { echo "  ${GREEN}✔${RESET}  $*"; }
info()     { echo "  ${CYAN}·${RESET}  $*"; }
warn()     { echo "  ${YELLOW}⚠${RESET}  $*"; }
err()      { echo "  ${RED}✘${RESET}  $*" >&2; }
die()      { err "$*"; exit 1; }
dim()      { echo "${DIM}$*${RESET}"; }

# Ask yes/no; default is supplied as 'y' or 'n'
confirm() {
  local msg="$1" default="${2:-y}"
  local prompt; [[ $default == y ]] && prompt="[Y/n]" || prompt="[y/N]"
  read -rp "     ${YELLOW}?${RESET}  $msg $prompt " ans
  ans="${ans:-$default}"
  [[ ${ans,,} == y ]]
}

# Numbered menu; sets $REPLY to chosen index (1-based)
pick_one() {
  local prompt="$1"; shift
  local options=("$@")
  echo "     ${YELLOW}?${RESET}  $prompt"
  for i in "${!options[@]}"; do
    printf "       %s) %s\n" "$((i+1))" "${options[$i]}"
  done
  while true; do
    read -rp "       Choice: " ans
    if [[ $ans =~ ^[0-9]+$ ]] && (( ans >= 1 && ans <= ${#options[@]} )); then
      REPLY=$ans; return
    fi
    echo "       Invalid choice, try again."
  done
}

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

# ── Dependency check ──────────────────────────────────────────────────────────
check_deps() {
  local missing=()
  for cmd in git stow pacman; do
    command -v "$cmd" &>/dev/null || missing+=("$cmd")
  done
  if (( ${#missing[@]} )); then
    die "Missing required tools: ${missing[*]}. Install them first."
  fi
  # Prefer paru, fall back to pacman -S for AUR-less installs
  if command -v paru &>/dev/null; then
    AUR_HELPER="paru"
  elif command -v yay &>/dev/null; then
    AUR_HELPER="yay"
  else
    AUR_HELPER=""
    warn "No AUR helper found (yay/paru). AUR packages will be skipped."
  fi
}

# ── Hardware detection ────────────────────────────────────────────────────────
detect_hardware_groups() {
  local groups=()
  # CPU
  if grep -qi intel /proc/cpuinfo 2>/dev/null; then
    groups+=("intel-cpu")
  elif grep -qi amd /proc/cpuinfo 2>/dev/null; then
    groups+=("amd-cpu")
  fi
  # GPU
  if lspci 2>/dev/null | grep -qi "intel.*graphics"; then
    groups+=("intel-gpu")
  fi
  if lspci 2>/dev/null | grep -qi "amd.*radeon\|advanced micro.*display"; then
    groups+=("amd-gpu")
  fi
  if lspci 2>/dev/null | grep -qi "nvidia"; then
    groups+=("nvidia")
  fi
  echo "${groups[@]}"
}

# ── Package-group helpers ─────────────────────────────────────────────────────

# Return the .pkgs file for a group name (strip suffixes variants)
pkg_file_for_group() {
  local group="$1"
  # Try exact match first, then with .pkgs extension
  local f
  for f in "$PACKAGES_DIR/${group}.pkgs" "$PACKAGES_DIR/${group}"; do
    [[ -f $f ]] && { echo "$f"; return; }
  done
}

# Read packages from a .pkgs file (skip comments & blanks)
read_pkg_file() {
  local file="$1"
  [[ -f $file ]] || return
  grep -v '^\s*#' "$file" | grep -v '^\s*$' | sort -u
}

# Load previously-saved list for a group
saved_pkg_list() {
  local group="$1"
  local f="$PKG_STATE_DIR/${group}"
  [[ -f $f ]] && cat "$f" || true
}

# Save current list for a group
save_pkg_list() {
  local group="$1"
  shift
  mkdir -p "$PKG_STATE_DIR"
  printf '%s\n' "$@" > "$PKG_STATE_DIR/${group}"
}

# Detect all available groups from packages/ dir
all_pkg_groups() {
  [[ -d $PACKAGES_DIR ]] || return
  find "$PACKAGES_DIR" -maxdepth 1 -name '*.pkgs' -printf '%f\n' \
    | sed 's/\.pkgs$//' | sort
}

# Read user-chosen optional groups from state
chosen_optional_groups() {
  state_get "optional_groups" | tr ',' '\n' | grep -v '^$' || true
}

# ── Git operations ────────────────────────────────────────────────────────────
git_fetch_and_check() {
  section "Checking for updates"
  cd "$DOTFILES_DIR"

  info "Fetching origin…"
  git fetch origin main --quiet

  local local_hash remote_hash
  local_hash=$(git rev-parse HEAD)
  remote_hash=$(git rev-parse origin/main)
  LAST_KNOWN_HASH=$(state_get "last_commit" || true)

  if [[ $local_hash == $remote_hash ]]; then
    ok "Already up to date (${local_hash:0:8})"
    HAS_UPDATE=false
    return
  fi

  info "New commits available:"
  git log --oneline "${local_hash}..${remote_hash}" | while read -r line; do
    dim "     $line"
  done
  HAS_UPDATE=true
  NEW_HASH="$remote_hash"
  OLD_HASH="$local_hash"
}

# ── Change detection ────────────────────────────────────────────────────────────────

# Files changed inside config/ (stow package) — used for conflict detection.
# Paths are stripped of the "config/" prefix so they match $HOME-relative targets.
collect_changed_config_files() {
  git -C "$DOTFILES_DIR" diff --name-only "${OLD_HASH}" "${NEW_HASH}" \
    | grep '^config/' \
    | sed 's|^config/||' \
    || true
}

# Returns 0 (true) if any file under packages/ changed between the two commits.
packages_changed() {
  git -C "$DOTFILES_DIR" diff --name-only "${OLD_HASH}" "${NEW_HASH}" \
    | grep -q '^packages/'
}

# Returns 0 (true) if any file under scripts/ changed between the two commits.
scripts_changed() {
  git -C "$DOTFILES_DIR" diff --name-only "${OLD_HASH}" "${NEW_HASH}" \
    | grep -q '^scripts/'
}

handle_conflicts() {
  local repo_changed_files
  mapfile -t repo_changed_files < <(collect_changed_config_files)

  if (( ${#repo_changed_files[@]} == 0 )); then
    return
  fi

  section "Checking for local conflicts"

  local conflicts=()
  for rel in "${repo_changed_files[@]}"; do
    # The stow target path
    local target="$STOW_TARGET/$rel"
    # Check if the symlink target (i.e. the real file in dotfiles) was
    # locally overridden — compare stowed hash vs. live hash
    if [[ -f $target && ! -L $target ]]; then
      # File exists but is NOT a symlink → user edited after un-stowing
      conflicts+=("$rel")
    elif [[ -L $target ]]; then
      # It's a symlink — compare live dotfiles version vs. what we last recorded
      local link_dest
      link_dest=$(readlink -f "$target" 2>/dev/null || true)
      local repo_file="$DOTFILES_DIR/$rel"
      # If the symlink points back into the repo, compare against old hash
      if [[ $link_dest == $repo_file ]]; then
        local old_blob new_blob
        old_blob=$(git -C "$DOTFILES_DIR" show "${OLD_HASH}:${rel}" 2>/dev/null \
                    | sha256sum | cut -d' ' -f1 || echo "")
        local live_blob
        live_blob=$(sha256sum "$link_dest" 2>/dev/null | cut -d' ' -f1 || echo "")
        # If live differs from old-repo → user modified through symlink
        if [[ -n $old_blob && $old_blob != $live_blob ]]; then
          conflicts+=("$rel")
        fi
      fi
    fi
  done

  if (( ${#conflicts[@]} == 0 )); then
    ok "No local conflicts"
    return
  fi

  warn "${#conflicts[@]} conflict(s) detected — files modified locally:"
  for f in "${conflicts[@]}"; do
    echo "       ${YELLOW}${f}${RESET}"
  done

  for rel in "${conflicts[@]}"; do
    echo
    info "Conflict: ${BOLD}${rel}${RESET}"
    pick_one "Which version to keep?" \
      "Upstream (repo) — discard local changes" \
      "Local — keep my version, skip this file"
    case $REPLY in
      1)
        ok "Will use upstream version of $rel"
        SKIP_FILES+=("__NONE__")  # placeholder so array stays parallel
        ;;
      2)
        warn "Skipping $rel — your local version will be kept"
        SKIP_FILES+=("$rel")
        ;;
    esac
  done
}

# ── Apply git update ──────────────────────────────────────────────────────────
apply_git_update() {
  section "Pulling updates"
  cd "$DOTFILES_DIR"

  # Stash any locally-modified tracked files so pull doesn't fail
  if ! git diff --quiet HEAD; then
    info "Stashing local modifications…"
    git stash push -m "updater-auto-stash" --quiet
    GIT_STASHED=true
  else
    GIT_STASHED=false
  fi

  git merge --ff-only origin/main --quiet
  ok "Repository updated to ${NEW_HASH:0:8}"

  # Re-apply stash if we created one (for files user chose to keep locally)
  if [[ $GIT_STASHED == true ]]; then
    info "Re-applying local stash…"
    git stash pop --quiet || {
      warn "Stash pop had conflicts — your stash is still saved (git stash list)"
    }
  fi
}

# ── Stow ─────────────────────────────────────────────────────────────────────
apply_stow() {
  section "Stowing dotfiles"

  local err_file; err_file=$(mktemp)
  if stow --dir="$STOW_DIR" --target="$STOW_TARGET" --restow "$STOW_PKG"        2>"$err_file"; then
    ok "stow: $STOW_PKG"
  else
    warn "stow: $STOW_PKG (errors below)"
    while IFS= read -r line; do
      echo "       ${DIM}${line}${RESET}"
    done < "$err_file"
  fi
  rm -f "$err_file"
}

# ── Package management ────────────────────────────────────────────────────────

# Install packages via pacman (official) or AUR helper
install_pkgs() {
  local pkgs=("$@")
  (( ${#pkgs[@]} == 0 )) && return
  if sudo pacman -S --needed --noconfirm "${pkgs[@]}" 2>/dev/null; then
    return
  fi
  # Some might be AUR-only — try AUR helper for the remainder
  if [[ -n $AUR_HELPER ]]; then
    "$AUR_HELPER" -S --needed --noconfirm "${pkgs[@]}" 2>/dev/null || true
  fi
}

# Remove packages (orphan-safe: only if nothing else depends on them)
remove_pkgs() {
  local pkgs=("$@")
  (( ${#pkgs[@]} == 0 )) && return
  # Filter to only actually-installed packages
  local installed=()
  for p in "${pkgs[@]}"; do
    pacman -Q "$p" &>/dev/null && installed+=("$p")
  done
  (( ${#installed[@]} == 0 )) && return
  sudo pacman -Rns --noconfirm "${installed[@]}" 2>/dev/null || true
}

process_pkg_group() {
  local group="$1"
  local pkg_file
  pkg_file=$(pkg_file_for_group "$group")

  if [[ -z $pkg_file ]]; then
    warn "No package file found for group '$group' — skipping"
    return
  fi

  mapfile -t new_pkgs < <(read_pkg_file "$pkg_file")
  mapfile -t old_pkgs < <(saved_pkg_list "$group")

  # Diff
  mapfile -t to_install < <(comm -23 \
    <(printf '%s\n' "${new_pkgs[@]}" | sort) \
    <(printf '%s\n' "${old_pkgs[@]+"${old_pkgs[@]}"}" | sort))
  mapfile -t to_remove < <(comm -13 \
    <(printf '%s\n' "${new_pkgs[@]}" | sort) \
    <(printf '%s\n' "${old_pkgs[@]+"${old_pkgs[@]}"}" | sort))

  local changed=false

  if (( ${#to_install[@]} )); then
    info "  [${group}] Installing: ${to_install[*]}"
    install_pkgs "${to_install[@]}" && changed=true
  fi

  if (( ${#to_remove[@]} )); then
    if confirm "  [${group}] Remove packages no longer in group: ${to_remove[*]}?" y; then
      remove_pkgs "${to_remove[@]}" && changed=true
    fi
  fi

  if [[ $changed == false ]] && (( ${#new_pkgs[@]} )); then
    ok "  [${group}] Up to date"
  fi

  # Always save current list (ensures state file exists after first run)
  save_pkg_list "$group" "${new_pkgs[@]}"
}

process_packages() {
  section "Package management"

  if [[ ! -d $PACKAGES_DIR ]]; then
    info "No packages/ directory found — skipping"
    return
  fi

  # 1. Auto groups (hardware-based)
  local hw_groups
  mapfile -t hw_groups < <(detect_hardware_groups)
  if (( ${#hw_groups[@]} )); then
    info "Detected hardware groups: ${hw_groups[*]}"
    for g in "${hw_groups[@]}"; do
      process_pkg_group "$g"
    done
  fi

  # 2. Core / always-on groups (everything that isn't hw-specific and not optional)
  #    Convention: groups NOT listed in optional.groups file are automatic
  local optional_file="$PACKAGES_DIR/optional.groups"
  local optional_group_names=()
  if [[ -f $optional_file ]]; then
    mapfile -t optional_group_names < <(grep -v '^\s*#' "$optional_file" | grep -v '^\s*$')
  fi

  local hw_pattern
  hw_pattern=$(printf '%s|' "${hw_groups[@]}" | sed 's/|$//')
  [[ -z $hw_pattern ]] && hw_pattern='__NOMATCH__'

  while IFS= read -r group; do
    # Skip hw groups already processed
    echo "$group" | grep -qE "^(${hw_pattern})$" && continue
    # Skip optional groups (handled below)
    local is_optional=false
    for og in "${optional_group_names[@]}"; do
      [[ $group == $og ]] && { is_optional=true; break; }
    done
    $is_optional && continue
    process_pkg_group "$group"
  done < <(all_pkg_groups)

  # 3. Optional groups — interactive first-time, then remembered
  if (( ${#optional_group_names[@]} )); then
    local chosen
    mapfile -t chosen < <(chosen_optional_groups)

    if (( ${#chosen[@]} == 0 )); then
      echo
      info "Optional package groups available. Choose which to install:"
      local selected=()
      for og in "${optional_group_names[@]}"; do
        if confirm "    Install group '${og}'?" n; then
          selected+=("$og")
        fi
      done
      if (( ${#selected[@]} )); then
        state_set "optional_groups" "$(IFS=,; echo "${selected[*]}")"
        chosen=("${selected[@]}")
      fi
    fi

    for og in "${chosen[@]}"; do
      process_pkg_group "$og"
    done
  fi
}

# ── AUR group (separate because needs AUR helper) ─────────────────────────────
process_aur_group() {
  [[ -z $AUR_HELPER ]] && { warn "Skipping AUR group (no AUR helper)"; return; }
  local pkg_file
  pkg_file=$(pkg_file_for_group "aur")
  [[ -z $pkg_file ]] && return

  section "AUR packages"
  mapfile -t new_pkgs < <(read_pkg_file "$pkg_file")
  mapfile -t old_pkgs < <(saved_pkg_list "aur")

  mapfile -t to_install < <(comm -23 \
    <(printf '%s\n' "${new_pkgs[@]}" | sort) \
    <(printf '%s\n' "${old_pkgs[@]+"${old_pkgs[@]}"}" | sort))
  mapfile -t to_remove < <(comm -13 \
    <(printf '%s\n' "${new_pkgs[@]}" | sort) \
    <(printf '%s\n' "${old_pkgs[@]+"${old_pkgs[@]}"}" | sort))

  if (( ${#to_install[@]} )); then
    info "Installing AUR: ${to_install[*]}"
    "$AUR_HELPER" -S --needed --noconfirm "${to_install[@]}" || true
  fi
  if (( ${#to_remove[@]} )); then
    if confirm "Remove AUR packages no longer listed: ${to_remove[*]}?" y; then
      remove_pkgs "${to_remove[@]}"
    fi
  fi
  (( ${#to_install[@]} == 0 && ${#to_remove[@]} == 0 )) && ok "AUR packages up to date"
  save_pkg_list "aur" "${new_pkgs[@]}"
}

# ── Entry point ───────────────────────────────────────────────────────────────
main() {
  echo
  echo "${BOLD}  Atlas Desktop Updater${RESET}"
  echo "${DIM}  Arch Linux${RESET}"
  echo "  ${DIM}Target: ${DOTFILES_DIR}${RESET}"

  [[ -d $DOTFILES_DIR/.git ]] || die "$DOTFILES_DIR is not a git repository."
  check_deps

  mkdir -p "$STATE_DIR" "$PKG_STATE_DIR"
  [[ -f $STATE_FILE ]] || touch "$STATE_FILE"

  HAS_UPDATE=false
  GIT_STASHED=false
  SKIP_FILES=()
  NEW_HASH=""
  OLD_HASH=""

  git_fetch_and_check

  if [[ $HAS_UPDATE == true ]]; then
    handle_conflicts
    apply_git_update

    # Re-stow config/ whenever dotfiles changed
    apply_stow

    # Re-stow scripts/ when it changed (keeps ~/.local/bin helpers up to date)
    if scripts_changed; then
      info "scripts/ changed — re-linking binaries"
      local err_file; err_file=$(mktemp)
      if stow --dir="$DOTFILES_DIR" --target="$HOME" --restow scripts \
           2>"$err_file"; then
        ok "stow: scripts"
      else
        warn "stow: scripts (errors below)"
        while IFS= read -r line; do
          echo "       ${DIM}${line}${RESET}"
        done < "$err_file"
      fi
      rm -f "$err_file"
    fi

    state_set "last_commit" "$NEW_HASH"

    # Re-process packages whenever packages/ changed, even if no git update
    # produced new dotfiles — package lists are independent of stowed files.
    if packages_changed; then
      info "packages/ changed — running full package sync"
    fi
  fi

  # Package sync runs on every invocation (diff vs saved state catches changes).
  # This also handles the case where a previous run partially failed.
  process_packages
  process_aur_group

  section "Done"
  ok "Dotfiles are up to date."
  if [[ $HAS_UPDATE == true ]]; then
    dim "     Updated: ${OLD_HASH:0:8} → ${NEW_HASH:0:8}"
  fi
  echo
}

main "$@"