#!/usr/bin/env bash
set -euo pipefail

HP_FILE="/etc/sysctl.d/99-performance-hugepages.conf"
MSR_FILE="/etc/modules-load.d/msr.conf"
CPUFREQ_FILE="/etc/default/cpufrequtils"

PACKAGES=(
  build-essential
  cmake
  git
  screen
  htop
  hwloc
  msr-tools
  numactl
)

say()  { echo -e "\n==> $*"; }
warn() { echo -e "!! $*"; }
ok()   { echo -e "-- $*"; }

need_sudo() {
  if [[ $EUID -ne 0 ]]; then
    command -v sudo >/dev/null || { echo "sudo not found"; exit 1; }
    SUDO=sudo
  else
    SUDO=""
  fi
}

ask() {
  local q="$1" d="${2:-y}" a
  local h="[y/N]"; [[ $d == y ]] && h="[Y/n]"
  while true; do
    read -r -p "$q $h: " a || true
    a="${a,,}"; [[ -z $a ]] && a="$d"
    case $a in
      y|yes) return 0 ;;
      n|no)  return 1 ;;
      *) echo "Введите y или n." ;;
    esac
  done
}

# ------------------ PACKAGES ------------------

pkg_installed() { dpkg -s "$1" >/dev/null 2>&1; }

install_packages() {
  say "ШАГ 1: базовые пакеты (compute/performance инструменты)"

  local missing=()
  for p in "${PACKAGES[@]}"; do
    pkg_installed "$p" || missing+=("$p")
  done

  if (( ${#missing[@]} == 0 )); then
    ok "Все нужные пакеты уже установлены. Пропускаем."
    return
  fi

  warn "Отсутствуют пакеты:"
  printf '  - %s\n' "${missing[@]}"

  if ask "Установить отсутствующие пакеты?" "y"; then
    $SUDO apt update
    $SUDO apt install -y "${missing[@]}"
    ok "Пакеты установлены."
  else
    warn "Установка пакетов пропущена."
  fi
}

# ------------------ DETECT ------------------

detect_hugepages() {
  [[ -f "$HP_FILE" ]] && sysctl -n vm.nr_hugepages 2>/dev/null || echo 0
}

detect_msr() {
  lsmod | grep -q '^msr' || [[ -f "$MSR_FILE" ]]
}

detect_governor() {
  cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor 2>/dev/null || echo unknown
}

any_perf_enabled() {
  [[ $(detect_hugepages) -gt 0 ]] || detect_msr || [[ $(detect_governor) == performance ]]
}

# ------------------ RAM + RECOMMEND ------------------

ram_total_gb() {
  awk '/MemTotal/ {printf "%.0f", $2/1024/1024}' /proc/meminfo
}

recommend_hp_gb() {
  local ram_gb
  ram_gb="$(ram_total_gb)"
  if (( ram_gb <= 16 )); then echo 2
  elif (( ram_gb <= 32 )); then echo 4
  else echo 8
  fi
}

# ------------------ HUGE PAGES ------------------

apply_hugepages() {
  local gb="$1"
  local pages=$(( gb * 512 )) # 1GB = 512 * 2MB pages
  echo "vm.nr_hugepages=$pages" | $SUDO tee "$HP_FILE" >/dev/null
  $SUDO sysctl --system >/dev/null
  ok "HugePages: ${gb} GB (${pages} pages) применено через $HP_FILE"
}

reset_hugepages() {
  $SUDO rm -f "$HP_FILE"
  $SUDO sysctl -w vm.nr_hugepages=0 >/dev/null
  ok "HugePages сброшены (vm.nr_hugepages=0, файл удалён)"
}

# ------------------ MSR ------------------

apply_msr() {
  $SUDO modprobe msr || true
  echo "msr" | $SUDO tee "$MSR_FILE" >/dev/null
  ok "MSR включён (модуль загружен, автозагрузка в $MSR_FILE)"
}

reset_msr() {
  $SUDO rm -f "$MSR_FILE"
  ok "MSR автозагрузка отключена (файл удалён). Модуль выгрузится после reboot (если не нужен другим)."
}

# ------------------ GOVERNOR ------------------

apply_governor() {
  if ! pkg_installed cpufrequtils; then
    warn "cpufrequtils не установлен."
    if ask "Установить cpufrequtils (нужно для governor=performance)?" "y"; then
      $SUDO apt update
      $SUDO apt install -y cpufrequtils
    else
      warn "Пропускаем governor."
      return
    fi
  fi

  echo 'GOVERNOR="performance"' | $SUDO tee "$CPUFREQ_FILE" >/dev/null
  $SUDO systemctl restart cpufrequtils 2>/dev/null || true
  ok "CPU governor = performance (файл $CPUFREQ_FILE)"
}

reset_governor() {
  $SUDO rm -f "$CPUFREQ_FILE"
  $SUDO systemctl restart cpufrequtils 2>/dev/null || true
  ok "CPU governor сброшен (файл $CPUFREQ_FILE удалён)"
}

# ------------------ STATUS ------------------

print_status() {
  say "ТЕКУЩИЙ СТАТУС"
  echo "RAM total:            $(ram_total_gb) GB"
  echo "HugePages configured: $(sysctl -n vm.nr_hugepages 2>/dev/null || echo unknown) pages"
  echo "HugePages file:       $HP_FILE $( [[ -f "$HP_FILE" ]] && echo '(exists)' || echo '(missing)' )"
  echo "MSR module loaded:    $(lsmod | grep -q '^msr' && echo yes || echo no)"
  echo "MSR autoload file:    $MSR_FILE $( [[ -f "$MSR_FILE" ]] && echo '(exists)' || echo '(missing)' )"
  echo "CPU governor:         $(detect_governor)"
  echo "Governor file:        $CPUFREQ_FILE $( [[ -f "$CPUFREQ_FILE" ]] && echo '(exists)' || echo '(missing)' )"
}

# ------------------ MAIN ------------------

main() {
  need_sudo
  say "Performance CPU Setup (интерактивный накат/откат настроек)"

  if ask "Сделать обновление системы (apt update + upgrade)?" "n"; then
    $SUDO apt update
    $SUDO apt upgrade -y
    ok "Обновление завершено."
  else
    ok "Обновление пропущено."
  fi

  install_packages

  print_status

  if any_perf_enabled; then
    warn "Обнаружены performance-настройки (возможно частично)."
    echo "1) Оставить как есть"
    echo "2) Привести настройки к умолчанию (откат)"
    echo "3) Перенастроить (сначала откат, потом накат)"
    while true; do
      read -rp "Выбор [1-3]: " c
      case "$c" in
        1) ok "Без изменений."; exit 0 ;;
        2) say "ОТКАТ НАСТРОЕК"; reset_hugepages; reset_msr; reset_governor; print_status; say "Готово. Рекомендуется reboot."; exit 0 ;;
        3) say "ПЕРЕНАСТРОЙКА: сначала откат"; reset_hugepages; reset_msr; reset_governor; break ;;
        *) echo "Введите 1, 2 или 3." ;;
      esac
    done
  else
    ok "Performance-настройки не обнаружены."
    if ! ask "Настроить систему под высокую CPU-производительность?" "y"; then
      exit 0
    fi
  fi

  # APPLY
  say "НАСТРОЙКА HugePages"
  local rec ramgb hp
  ramgb="$(ram_total_gb)"
  rec="$(recommend_hp_gb)"
  ok "ОЗУ: ${ramgb} GB. Рекомендация HugePages: ${rec} GB."
  read -rp "Введите размер HugePages в ГБ (0 = отключить) [по умолчанию $rec]: " hp
  hp="${hp:-$rec}"

  if ! [[ "$hp" =~ ^[0-9]+$ ]]; then
    warn "Некорректное число: $hp. Использую рекомендованное: $rec."
    hp="$rec"
  fi

  if (( hp == 0 )); then
    reset_hugepages
  else
    apply_hugepages "$hp"
  fi

  say "НАСТРОЙКА MSR"
  if ask "Включить MSR модуль и автозагрузку?" "y"; then
    apply_msr
  else
    warn "MSR пропущен."
  fi

  say "НАСТРОЙКА CPU governor"
  if ask "Установить governor=performance?" "y"; then
    apply_governor
  else
    warn "Governor пропущен."
  fi

  print_status
  say "Готово. Рекомендуется перезагрузка."
  if ask "Перезагрузить сейчас?" "n"; then
    $SUDO reboot
  fi
}

main "$@"
