#!/usr/bin/env bash

set -o pipefail

# echo wrappers
INFO(){ echo "INFO: $*";}
WARN(){ echo "WARN: $*";}
ERRO(){ echo "ERRO: $*"; exit 1;}
INFO_E(){ echo "INFO: $*"; exit 0;}

# Helpers
YN(){
  l="${1,,}"
  [[ "${l}" =~ (yes|y|1|true) ]] && return 0
  return 1
}

write(){
  { [ -z "$1" ] || [ -z "$2" ]; } && return
  DATA="$1" FILE="$2"
  echo "${DATA}" > "${FILE}"
}

# files with one byte data will use one page of RAM in tmpfs,
# i.e. 4KiB on x86_64
# to avoid that, store short strings in symbolic links
write_l(){
  { [ -z "$1" ] || [ -z "$2" ]; } && return
  DATA="$1" FILE="$2"
  DATA=$(echo "${DATA}" | base64)
  ln -sfr "/${DATA}" "${FILE}"
}

read_l(){
  [ -z "$1" ] && return
  FILE="$1"
  DATA=$(basename "$(realpath "${FILE}")")
  echo "${DATA}" | base64 -d
}

get_mem_stat_multi() {
  # expects name of output array as $1, followed by n wanted keys
  local -ar fields=("${@:2}")
  local -n outref="$1"
  local -i matches=0
  while read -r line; do
    for field in "${fields[@]}"; do
      if [[ $line =~ ^$field:[[:space:]]+([[:digit:]]+)[[:space:]]+kB$ ]]; then
          outref+=( $(( BASH_REMATCH[1] * 1024 )) )
          ((++matches))
      fi
    done
    ((matches == ${#fields[@]})) && break;
  done < /proc/meminfo
}

get_mem_stat() {
  local -a value
  get_mem_stat_multi value "$1"
  printf "%s" "${value[0]}"
}

get_fs_type(){ df "$1" --output=fstype | tail -n 1; }
AMI_ROOT(){ [ "${UID}" == "0" ] || ERRO "Script must be run as root!"; }
FIND_SWAP_UNITS(){ find /run/systemd/system/ /run/systemd/generator/ -type f -name "*.swap"; }
GET_WHAT_FROM_SWAP_UNIT(){ grep -oP 'What=\K.*' "$1"; }
help(){
  echo "$0 start|stop|status"
  echo "   start  - init daemon"
  echo "   stop   - stop daemon"
  echo "   status - show some swap status info"
}

snore()
{
    local IFS
    [[ -n "${_snore_fd:-}" ]] || exec {_snore_fd}<> <(:)
    read ${1:+-t "$1"} -u $_snore_fd || :
}

# Global vars
readonly RUN_SYSD="/run/systemd"
readonly NCPU=$(nproc)
readonly RAM_SIZE=$(get_mem_stat MemTotal)
readonly PAGE_SIZE=$(getconf PAGESIZE)
readonly ETC_SYSD="/etc/systemd"
readonly CONFIG="${ETC_SYSD}/swap.conf"
readonly WORK_DIR="/run/systemd/swap"

readonly LOCK_STARTED="${WORK_DIR}/.started"
readonly ZSWAP_M="/sys/module/zswap"
readonly ZSWAP_M_P="/sys/module/zswap/parameters"

readonly KMAJOR=$(uname -r | cut -d'.' -f1)
#readonly KMINOR=$(uname -r | cut -d'.' -f2)

# swap unit generator
gen_swap_unit(){
  export What Priority Options Tag Type
  for i in "$@"; do
    case $i in
      What=*)
        What="${i//What=/}"
      ;;
      Priority=*)
        Priority="${i//Priority=/}"
      ;;
      Options=*)
        Options="${i//Options=/}"
      ;;
      Tag=*)
        Tag="${i//Tag=}"
      ;;
    esac
  done
  [ -n "${What}" ] || return 1

  What="$(realpath "${What}")"
  # assume it's a file by default
  Type="File"
  if [ -b "${What}" ]; then
    Type="Block/Partition"
    [[ "${What}" =~ loop ]] && Type="File"
  fi

  UNIT_NAME=$(systemd-escape -p "${What}")
  UNIT_PATH="${RUN_SYSD}/system/${UNIT_NAME}.swap"
  {
    echo '[Unit]'
    echo "Description=Swap ${Type}"
    echo 'Documentation=https://github.com/Nefelim4ag/systemd-swap'
    echo
    echo "# Generated by systemd-swap"
    echo "# Tag=${Tag}"
    echo
    echo '[Swap]'
    echo "What=${What}"
    echo "TimeoutSec=1h"
    [ -z "${Priority}" ] || echo "Priority=${Priority}"
    [ -z "${Options}"  ] || echo "Options=${Options}"
  } > "${UNIT_PATH}"

  ln -srf "${UNIT_PATH}" "${RUN_SYSD}/system/swap.target.wants"

  if [[ "${Type}" == "File" ]]; then
    ln -srf "${UNIT_PATH}" "${RUN_SYSD}/system/local-fs.target.wants"
  fi

  echo "${UNIT_NAME}.swap"
}

# INIT

AMI_ROOT

mkdir -p \
  "${WORK_DIR}" \
  "${RUN_SYSD}/system/local-fs.target.wants" \
  "${RUN_SYSD}/system/swap.target.wants"

case "$1" in
  start)
    [ -f "${LOCK_STARTED}" ] && ERRO "$0 already started"
    touch "${LOCK_STARTED}"

    INFO "Load: ${CONFIG}"
    # shellcheck source=swap.conf
    . "${CONFIG}" || ERRO "Error loading ${CONFIG}"

    declare -A CONFS
    for L_CONF in {/lib,/run,/etc}/systemd/swap.conf.d/*.conf; do
      [[ ! -r ${L_CONF} ]] && continue
      BASE=$(basename "${L_CONF}")
      [[ ${CONFS[${BASE}]} ]] && BASES=("${BASES[@]/${BASE}}")
      CONFS["${BASE}"]="${L_CONF}"
      BASES+=("${BASE}")
    done
    for BASE in "${BASES[@]}"; do
      [[ ! ${BASE} ]] && continue
      L_CONF="${CONFS[${BASE}]}"
      # shellcheck source=swap.conf
      . "${L_CONF}" || ERRO "Error loading ${L_CONF}"
    done

    zswap_enabled=${zswap_enabled:-0}
    if YN "${zswap_enabled}"; then
      [ ! -d ${ZSWAP_M} ] && ERRO "Zswap - not supported on current kernel"

      zswap_compressor=${zswap_compressor:-zstd}
      zswap_max_pool_percent=${zswap_max_pool_percent:-25}
      zswap_zpool=${zswap_zpool:-z3fold}

      INFO "Zswap: backup current configuration: start"
      mkdir -p "${WORK_DIR}/zswap"
      for file in "${ZSWAP_M_P}"/*; do
        read -r VAL < "${file}"
        write_l "${VAL}" "${WORK_DIR}/zswap/$(basename "${file}")"
      done
      INFO "Zswap: backup current configuration: complete"
      INFO "Zswap: set new parameters: start"
      INFO "Zswap: Enable: ${zswap_enabled}, Comp: ${zswap_compressor},  Max pool %: ${zswap_max_pool_percent}, Zpool: ${zswap_zpool}"
      write "${zswap_enabled}"          "${ZSWAP_M_P}"/enabled
      write "${zswap_compressor}"       "${ZSWAP_M_P}"/compressor
      write "${zswap_max_pool_percent}" "${ZSWAP_M_P}"/max_pool_percent
      write "${zswap_zpool}"            "${ZSWAP_M_P}"/zpool
      INFO "Zswap: set new parameters: complete"
    fi

    zram_enabled=${zram_enabled:-0}
    if YN "${zram_enabled}"; then
      [ -z "${zram_size}" ] && zram_size=$(( RAM_SIZE / 4 ))
      zram_count=${zram_count:-${NCPU}}
      zram_size=$(( zram_size / zram_count ))
      zram_streams=${zram_streams:-${NCPU}}
      zram_alg=${zram_alg:-"zstd"}
      zram_prio=${zram_prio:-"32767"}
      zram_dev=""

      INFO "Zram: check availability"
      if [ ! -d "/sys/module/zram" ]; then
        INFO "Zram: not part of kernel, trying to find zram module"
        modprobe -n zram || ERRO "Zram: can't find zram module!"
        # workaround for some zram initialization problems
        for (( i = 0; i < 10; i++ )); do
          [ -d "/sys/module/zram" ] && break
          modprobe zram
          snore 1
        done
        INFO "Zram: module successfully loaded"
      else
        INFO "Zram: module already loaded"
      fi

      for (( i = 0; i < zram_count; i++ )); do
        INFO "Zram: trying to initialize free device"
        # zramctl is an external program -> return path to first free device
        OUTPUT=$(zramctl -f -a "${zram_alg}" -t "${zram_streams}" -s "${zram_size}" 2>&1)
        case "${OUTPUT}" in
          *"failed to reset: Device or resource busy"*)
            snore 1
          ;;
          *"zramctl: no free zram device found"*)
            WARN "Zram: zramctl can't find free device"
            INFO "Zram: using workaround hook for hot add"
            [ ! -f /sys/class/zram-control/hot_add ] && \
              ERRO "Zram: this kernel does not support hot adding zram devices, please use a 4.2+ kernel or see 'modinfo zram´ and create a modprobe rule"
            read -r NEW_ZRAM < /sys/class/zram-control/hot_add
            INFO "Zram: success: new device /dev/zram${NEW_ZRAM}"
          ;;
          /dev/zram*)
            [ -b "${OUTPUT}" ] || continue
            zram_dev="${OUTPUT}"
          ;;
        esac
        if [ -b "${zram_dev}" ]; then
          INFO "Zram: initialized: ${zram_dev}"
          mkswap "${zram_dev}" &> /dev/null && \
          UNIT_NAME=$(gen_swap_unit What="${zram_dev}" Options=discard Priority="${zram_prio}" Tag=zram)
          systemctl daemon-reload
          systemctl start "${UNIT_NAME}"
        else
          WARN "Can't get free Zram device"
        fi
      done
    fi

    swapfc_enabled=${swapfc_enabled:-0}
    if YN "${swapfc_enabled}"; then
      chunk_size=${swapfc_chunk_size:-"512M"}
      swapfc_max_count=${swapfc_max_count:-"8"}
      swapfc_path=${swapfc_path:-"/var/lib/systemd-swap/swapfc/"}
      swapfc_frequency=${swapfc_frequency:-"1"}
      swapfc_force_use_loop=${swapfc_force_use_loop:-"0"}
      swapfc_nocow=${swapfc_nocow:-"1"}
      swapfc_directio=${swapfc_directio:-"1"}
      swapfc_force_preallocated=${swapfc_force_preallocated:-"0"}
      swapfc_free_swap_perc=${swapfc_free_swap_perc:-"15"}

      get_free_swap_perc(){
        local -a swap_stats
        get_mem_stat_multi swap_stats SwapTotal SwapFree
        SWAP_USED=$(( swap_stats[0] - swap_stats[1] ))
        # minimum for total is 1 to prevent divide by zero
        echo $(( (swap_stats[1] * 100) / (swap_stats[0] + 1) ));
      }
      
      get_free_ram_perc(){
        local -a ram_stats
        get_mem_stat_multi ram_stats MemTotal MemFree
        echo $(( (ram_stats[1] * 100) / ram_stats[0] ));
      }

      to_bytes(){ numfmt --to=none --from=iec "$1"; }

      # Must exists before stat()
      mkdir -p "${swapfc_path}"

      chunk_size=$(to_bytes "${chunk_size}")
      BLOCK_SIZE=$(stat -f -c %s "${swapfc_path}")
      FSTYPE=$(get_fs_type "${swapfc_path}")

      if [ "$FSTYPE" = "btrfs" ]; then
        # if btrfs supports regular swap files(kernel version 5+), force disable COW to avoid data corruption
        # if it doesn't, use the old swap through loop workaround
        if [ "${KMAJOR}" -ge 5 ]; then
              swapfc_nocow=true
        else
              FSTYPE="btrfs_old"
        fi
      fi

      YN "${swapfc_force_use_loop}" && FSTYPE=btrfs_old
      [ "${FSTYPE}" != "btrfs_old" ] && swapfc_force_preallocated=1

      check_ENOSPC(){
        path="$1"
        # check free space for avoiding problems on swap io + ENOSPC
        FREE_BLOCKS=$(stat -f -c %f "${path}" )
        FREE_BYTES=$(( FREE_BLOCKS * BLOCK_SIZE ))
        # also try leaving some free space
        FREE_BYTES=$(( FREE_BYTES - chunk_size ))
        ((FREE_BYTES < chunk_size)) && return 0
        return 1
      }

      prepare_swapfile(){
        chunk_size="$1" file="$2"
        touch "${file}" && chmod 0600 "${file}"

        [ "$FSTYPE" = "btrfs" ] && YN "${swapfc_nocow}" && chattr +C "${file}"
        if [[ $FSTYPE = ext4 || $FSTYPE = ext3 || $FSTYPE = xfs || $FSTYPE = f2fs ]]; then
          dd if=/dev/zero of="$file" status=none bs=1M count=$(( chunk_size / 1048576))
        else
          if YN ${swapfc_force_preallocated}; then
            # ext2/3 does not support fallocate
            fallocate -l "${chunk_size}" "${file}" \
            || truncate -s "${chunk_size}" "${file}"
          else
            truncate -s "${chunk_size}" "${file}"
          fi
        fi

        losetup_w(){
          file="$1"
          if YN "${swapfc_directio}"; then
            losetup -f --show --direct-io=on  "${file}"
          else
            losetup -f --show --direct-io=off "${file}"
          fi
          # loop uses a file descriptor - if the file still exists,
          # but does not have a path like O_TMPFILE
          # when loop detaches a file, the file will be deleted.
          rm "${file}"
        }

        RET="${file}"
        case "${FSTYPE}" in
          ext2|ext3|ext4|btrfs)
          ;;
          btrfs_old)
            RET=$(losetup_w "${file}")
          ;;
          *)
            # -z would result in two overwrites, so explicitly use
            # /dev/zero as source for a quicker *single* iteration
            shred -n1 --random-source=/dev/zero "${file}"
          ;;
        esac
        echo "${RET}"
      }

      if (( "${swapfc_max_count}" > 32 )) || (( 1 > "${swapfc_max_count}" )); then
        WARN "swapfc_max_count must be in range 1..32, reset to 1"
        swapfc_max_count=1
      fi

      mkdir -p \
        "${swapfc_path}" \
        "${WORK_DIR}/swapfc/"
      touch "${WORK_DIR}/swapfc/.lock"
      {
        allocated=0
        free_threshold_up=${swapfc_free_swap_perc}
        free_threshold_down=$(( swapfc_free_swap_perc + 40 ))
        while [ -f "${WORK_DIR}/swapfc/.lock" ] && snore "${swapfc_frequency}"; do
          if [ $allocated -eq 0 ]; then
            curr_free_ram_perc=$(get_free_ram_perc)
            if (( curr_free_ram_perc < free_threshold_up)); then
              if check_ENOSPC "${swapfc_path}"; then
                WARN "swapFC: ENOSPC"
                continue
              fi
              allocated=$(( allocated + 1 ))
              INFO "swapFC: free ram: ${curr_free_ram_perc} < ${free_threshold_up} - allocate chunk: ${allocated}"
              swapfile=$(prepare_swapfile "${chunk_size}" "${swapfc_path}/${allocated}")
              mkswap -L SWAP_${FSTYPE}_${allocated} "${swapfile}" | sed ':a;N;$!ba;s/)\n/), /g'
              if YN ${swapfc_force_preallocated}; then
                UNIT_NAME=$(gen_swap_unit What="${swapfile}" Tag="swapfc_${allocated}")
              else
                UNIT_NAME=$(gen_swap_unit What="${swapfile}" Options=discard Tag="swapfc_${allocated}")
              fi
              systemctl daemon-reload
              systemctl start "${UNIT_NAME}"
              if [ -b "${swapfile}" ]; then
                losetup -d "${swapfile}"
              fi
              continue
            else
              continue
            fi
          fi
          curr_free_swap_perc=$(get_free_swap_perc)
          if (( curr_free_swap_perc < free_threshold_up )) && (( allocated < swapfc_max_count )); then
            if check_ENOSPC "${swapfc_path}"; then
              WARN "swapFC: ENOSPC"
              continue
            fi
            allocated=$(( allocated + 1 ))
            INFO "swapFC: free swap: ${curr_free_swap_perc} < ${free_threshold_up} - allocate chunk: ${allocated}"
            swapfile=$(prepare_swapfile "${chunk_size}" "${swapfc_path}/${allocated}")
            mkswap -L SWAP_${FSTYPE}_${allocated} "${swapfile}" | sed ':a;N;$!ba;s/)\n/), /g'
            if YN ${swapfc_force_preallocated}; then
              UNIT_NAME=$(gen_swap_unit What="${swapfile}" Tag="swapfc_${allocated}")
            else
              UNIT_NAME=$(gen_swap_unit What="${swapfile}" Options=discard Tag="swapfc_${allocated}")
            fi
            systemctl daemon-reload
            systemctl start "${UNIT_NAME}"
            if [ -b "${swapfile}" ]; then
              losetup -d "${swapfile}"
            fi
          fi
          (( allocated <= 2 )) && continue
          if (( curr_free_swap_perc > free_threshold_down )); then
            INFO "swapFC: free swap: ${curr_free_swap_perc} > ${free_threshold_down} - freeup chunk: ${allocated}"
            FIND_SWAP_UNITS | while read -r UNIT_PATH; do
              if grep -q swapfc_${allocated} "${UNIT_PATH}"; then
                DEV=$(GET_WHAT_FROM_SWAP_UNIT "${UNIT_PATH}")
                UNIT_NAME=$(basename "${UNIT_PATH}")
                systemctl stop "${UNIT_NAME}" || swapoff "${DEV}"
                rm -vf "${UNIT_PATH}"
                [ -f "${DEV}" ] && rm -f "${DEV}"
              fi
            done
            allocated=$(( allocated - 1 ))
          fi
        done
      } &
    fi

    swapd_auto_swapon=${swapd_auto_swapon:-1}
    swapd_prio=${swapd_prio:-1024}
    if YN "${swapd_auto_swapon}"; then
      INFO "swapD: pickup devices from systemd-gpt-auto-generator"
      FIND_SWAP_UNITS | while read -r UNIT_PATH; do
        if grep -q systemd-gpt-auto-generator "${UNIT_PATH}"; then
          DEV=$(GET_WHAT_FROM_SWAP_UNIT "${UNIT_PATH}")
          UNIT_NAME=$(basename "${UNIT_PATH}")
          swapoff "${DEV}"
          rm -vf "${UNIT_PATH}"
        fi
      done

      INFO "swapD: searching swap devices"
      mkdir -p ${WORK_DIR}/swapd/
      for device in $(blkid -t TYPE=swap -o device | grep -vE '(zram|loop)'); do
        for used_device in $(swapon --show=NAME --noheadings); do
          [ "${device}" == "${used_device}" ] && unset device
        done
        [ ! -b "${device}" ] && continue
        UNIT_NAME=$(gen_swap_unit What="${device}" Options=discard Tag="swapd")
        systemctl daemon-reload
        systemctl start "${UNIT_NAME}" || continue
        INFO "swapD: enabled device: ${device}"
        swapd_prio=$(( swapd_prio - 1 ))
      done
    fi
  ;;
  stop)
    [ ! -f "${LOCK_STARTED}" ] && INFO_E "$0 already stopped"

    FIND_SWAP_UNITS | while read -r UNIT_PATH; do
      if grep -q swapd "${UNIT_PATH}"; then
        DEV=$(GET_WHAT_FROM_SWAP_UNIT "${UNIT_PATH}")
        UNIT_NAME=$(basename "${UNIT_PATH}")
        swapoff "${DEV}"
        rm -vf "${UNIT_PATH}"
      fi
    done

    if [ -f "${WORK_DIR}/swapfc/.lock" ]; then
      rm -vf "${WORK_DIR}/swapfc/.lock"
      snore 5
    fi

    FIND_SWAP_UNITS | while read -r UNIT_PATH; do
      if grep -q swapfc "${UNIT_PATH}"; then
        DEV=$(GET_WHAT_FROM_SWAP_UNIT "${UNIT_PATH}")
        UNIT_NAME=$(basename "${UNIT_PATH}")
        swapoff "${DEV}"
        rm -vf "${UNIT_PATH}"
        [ -f "${DEV}" ] && rm -f "${DEV}"
      fi
    done

    FIND_SWAP_UNITS | while read -r UNIT_PATH; do
      if grep -q zram "${UNIT_PATH}"; then
        DEV=$(GET_WHAT_FROM_SWAP_UNIT "${UNIT_PATH}")
        UNIT_NAME=$(basename "${UNIT_PATH}")
        INFO "Zram: swapoff ${DEV}"
        swapoff "${DEV}"
        rm -vf "${UNIT_PATH}"
        zramctl -r "${DEV}"
      fi
    done

    if [ -d "${WORK_DIR}/zswap" ]; then
      INFO "Zswap: restore configuration: start"
      for file in "${WORK_DIR}/zswap"/*; do
        read_l "${file}" > "${ZSWAP_M_P}/$(basename "${file}")"
      done
      INFO "Zswap: restore configuration: complete"
    fi

    wait

    rm -vf "${LOCK_STARTED}"
  ;;
  status)
    declare -a swap_stats
    get_mem_stat_multi swap_stats SwapTotal SwapFree
    SWAP_USED=$(( swap_stats[0] - swap_stats[1] ))
    unset swap_stats

    if [ -d /sys/module/zswap ]; then
      echo Zswap:
      read -r USED_BYTES < /sys/kernel/debug/zswap/pool_total_size
      USED_PAGES=$(( USED_BYTES / PAGE_SIZE ))
      read -r STORED_PAGES < /sys/kernel/debug/zswap/stored_pages
      STORED_BYTES=$(( STORED_PAGES * PAGE_SIZE ))
      RATIO=0
      (( STORED_PAGES > 0 )) && RATIO=$(( USED_PAGES * 100 / STORED_PAGES ))
      {
        grep . /sys/module/zswap/parameters/* | \
          cut -d'/' -f 6 | \
          sed -e 's/:/ /g' -e 's/^/. /g'
      } | column -t
      {
        grep . /sys/kernel/debug/zswap/* | \
          cut -d'/' -f6 | \
          sed -e 's/:/ /g' -e 's/^/. . /g' | column -t
        echo . . compress_ratio ${RATIO}%
        (( SWAP_USED > 0 )) && \
          echo . . zswap_store/swap_store ${STORED_BYTES}/${SWAP_USED} $(( STORED_BYTES * 100 / SWAP_USED ))%
      } | column -t

    else
      WARN "Zswap unavailable"
    fi

    if zramctl | grep -q '\[SWAP\]'; then
      echo ZRam:
      zramctl | \
        grep -e '^NAME\|\[SWAP\]' | \
        sed -e 's/\(MOUNTPOINT\|\[SWAP\]\)$//g' -e 's/^/.\t/g' | \
        column -t | \
        uniq
    fi

    if [ -d /run/systemd/swap/swapd ]; then
      echo swapD:
      swapon --raw --show TYPE | \
        grep -v 'zram\|file\|loop' | \
        sed 's/^/.\t/g' | \
        column -t
    fi

    if [ -d /run/systemd/swap/swapfc ]; then
      echo swapFC:
      swapon --raw --show TYPE | \
        grep -e 'NAME\|file\|loop' | \
        sed 's/^/.\t/g' | \
        column -t
    fi
  ;;
  *)
    help
  ;;
esac
