#!/usr/bin/env bash

set -euo pipefail
shopt -s inherit_errexit
IFS=$'\n\t'

# =============================================================================
# Script: add_swap.sh
# Description: Linux 添加 Swap 文件通用脚本（标准文档推荐方式）
# Version: 2.0.0
# Compatibility: Bash 5.0+
# =============================================================================

readonly SCRIPT_VERSION="2.0.0"
readonly DEFAULT_SWAP_SIZE="4G"
readonly SWAPFILE="/swapfile"

# ----------------------------------------------------------------------------
# 函数: utils::log
# 描述: 统一日志输出到 stderr
# ----------------------------------------------------------------------------
utils::log() {
  local level="$1"
  local msg="$2"
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] [${level}] ${msg}" >&2
}

# ----------------------------------------------------------------------------
# 函数: check_root
# 描述: 检查是否以 root 权限运行
# ----------------------------------------------------------------------------
check_root() {
  if [[ "${EUID}" -ne 0 ]]; then
    utils::log "ERROR" "此脚本必须以 root 或 sudo 权限运行，请使用: sudo $0"
    exit 1
  fi
}

# ----------------------------------------------------------------------------
# 函数: create_swap_file
# 描述: 创建 swap 文件并立即启用（优先 fallocate，回退 dd）
# ----------------------------------------------------------------------------
create_swap_file() {
  local swap_size="$1"
  local swapfile="$2"

  utils::log "INFO" "正在创建 ${swap_size} 的 swap 文件 → ${swapfile}"

  # 如果文件已存在，先清理
  if [[ -f "${swapfile}" ]]; then
    if swapon --show | grep -q "${swapfile}"; then
      utils::log "INFO" "检测到已有 swap 正在使用，先禁用..."
      swapoff "${swapfile}" || true
    fi
    rm -f "${swapfile}"
  fi

  # 优先使用 fallocate（更快、更安全）
  if ! fallocate -l "${swap_size}" "${swapfile}" 2>/dev/null; then
    utils::log "WARNING" "fallocate 不可用或失败，回退使用 dd 创建..."
    # 纯 Bash 解析大小为 MB（无需 bc/numfmt/awk，完全兼容 Bash 5.0+）
    local clean_size="${swap_size^^}"
    local num="${clean_size%%[GM]*}"
    local mb=4096
    if [[ "${clean_size}" == *G* ]]; then
      mb=$((num * 1024))
    elif [[ "${clean_size}" == *M* ]]; then
      mb="${num}"
    fi
    dd if=/dev/zero of="${swapfile}" bs=1M count="${mb}" status=progress
  fi

  chmod 600 "${swapfile}"
  mkswap "${swapfile}"
  swapon "${swapfile}"

  utils::log "SUCCESS" "Swap 文件创建并启用成功"
}

# ----------------------------------------------------------------------------
# 函数: add_to_fstab
# 描述: 永久添加到 /etc/fstab（幂等）
# ----------------------------------------------------------------------------
add_to_fstab() {
  local swapfile="$1"

  if ! grep -q "^${swapfile}[[:space:]]" /etc/fstab 2>/dev/null; then
    echo "${swapfile} none swap sw 0 0" >> /etc/fstab
    utils::log "INFO" "已添加到 /etc/fstab（开机自动挂载）"
  else
    utils::log "INFO" "fstab 中已存在该 swap 记录"
  fi
}

# ----------------------------------------------------------------------------
# 主函数
# ----------------------------------------------------------------------------
main() {
  check_root

  local swap_size="${1:-${DEFAULT_SWAP_SIZE}}"
  local swapfile="${SWAPFILE}"

  # 如果未通过参数传入，则交互式询问
  if [[ -z "${1:-}" ]]; then
    read -r -p "请输入要添加的 Swap 大小（例如 4G、8G、16G，默认 ${DEFAULT_SWAP_SIZE}）: " user_input
    swap_size="${user_input:-${DEFAULT_SWAP_SIZE}}"
  fi

  # 简单校验大小格式（允许 数字 + G/M/g/m）
  if [[ ! "${swap_size}" =~ ^[0-9]+[GgMm]?$ ]]; then
    utils::log "ERROR" "大小格式错误！请使用如 4G、8G、4096M"
    exit 1
  fi

  create_swap_file "${swap_size}" "${swapfile}"
  add_to_fstab "${swapfile}"

  utils::log "INFO" "=== 当前内存与 Swap 状态 ==="
  free -h
  swapon --show
}

main "$@"