#!/usr/bin/env bash
# Zmount automatically configures zram devices with the ztab config file
# Copyright (C)  2021  Aaron Ball <nullspoon@oper.io>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

set -euo pipefail
export IFS=$'\n\t'

CONFIG='/etc/ztab'
RUN='/var/run/zmount'

_parseline() {
  local -a lineargs=()
  lineargs=($(printf '%s\n' "${1:?}" | tr ' ' '\n'))
  if [ "${#lineargs[@]}" -ne 5 ]; then
    printf 'ERROR: Line "%s" does not have correct columns\n' >&2
    return 1
  fi
  printf '%s\n' ${lineargs[@]}
  return 0
}

_setperms() {
  local dir="${1}"
  local permstr="${2}"
  local perm
  local -a keyval=()

  [ "${permstr}" = 'none' ] && return 0

  for perm in $(printf '%s\n' ${permstr} | tr ',' '\n'); do
    keyval=($(printf '%s\n' "${perm}" | tr '=' '\n'))
    if [ "${keyval[0]}" = 'user' ]; then
      chown "${keyval[1]}" "${dir}"
    elif [ "${keyval[0]}" = 'group' ]; then
      chgrp "${keyval[1]}" "${dir}"
    elif [ "${keyval[0]}" = 'mode' ]; then
      chmod "${keyval[1]}" "${dir}"
    fi
  done
}

_setup_dev() {
  local fs="${1}"
  local size="${2}"
  local dir="${3}"
  local opts="${4:-defaults}"
  local perms="${5:-none}"      # Default to 777 to emulate tmpfs

  local dev="$(zramctl -f)"
  local threads="$(( $(nproc) / 2 ))"

  [ ! -d "${RUN}" ] && mkdir -p ${RUN}

  # Check for existing mounts to avoid duplicates
  if [ "${fs}" = 'swap' ] && [ -f "${RUN}/swap" ]; then
    printf 'ERROR: Zswap device already detected. Skipping second.\n' >&2
    return 1
  elif [ -f "${RUN}/${dir//\//_}" ]; then
    printf 'ERROR: Skipping duplicate mount dir %s\n' "${dir}" >&2
    return 1
  fi

  # Create the zram device
  zramctl "${dev}" -s "${size}" -t "${threads}"

  if [ "${fs}" = 'swap' ]; then
    mkswap "${dev}" 2>&1 1>/dev/null
    swapon "${dev}"
    printf '%s\n' "${dev}" > "${RUN}/swap"
  else
    if [ "${fs}" = 'btrfs' ]; then
      mkfs.btrfs -q "${dev}"
    elif [ "${fs}" = 'ext2' ]; then
      mkfs.ext2 -q "${dev}"
    elif [ "${fs}" = 'ext3' ]; then
      mkfs.ext3 -q "${dev}"
    elif [ "${fs}" = 'ext4' ]; then
      mkfs.ext4 -q "${dev}"
    elif [ "${fs}" = 'exfat' ]; then
      mkfs.exfat -q "${dev}"
    elif [ "${fs}" = 'vfat' ]; then
      mkfs.vfat -q "${dev}"
    elif [ "${fs}" = 'xfs' ]; then
      mkfs.xfs -q "${dev}"
    else
      printf 'Unknown filesystem type %s\n' "${fs}"
      return 1
    fi
    mount -o "${opts}" "${dev}" "${dir}"
    printf '%s\n' "${dev}" > "${RUN}/${dir//\//_}"
  fi
}

status() { zramctl; }

start() {
  local -a lineargs=()

  for line in $(grep -v '^#' /etc/ztab); do
    lineargs=($(_parseline "${line}"))

    if [ "${lineargs[0]}" = 'swap' ]; then
      printf 'Setting up %s zswap\n' "${lineargs[1]}"
    else
      printf 'Setting up %s zram dev on %s\n' "${lineargs[1]}" "${lineargs[2]}"
    fi
    _setup_dev ${lineargs[@]}
    _setperms "${lineargs[2]}" "${lineargs[4]}"
  done
}


stop() {
  local -a lineargs=()
  local fs size dir

  for line in $(grep -v '^#' /etc/ztab); do
    lineargs=($(_parseline "${line}"))
    fs="${lineargs[0]}"
    size="${lineargs[1]}"
    dir="${lineargs[2]}"
    # Ignore opts and perms column

    if [ "${fs}" = 'swap' ] && [ -f "${RUN}/swap" ]; then
      dev="$(cat ${RUN}/swap)"
      printf 'Stopping %s (mounted as swap)\n' "${dev}"
      swapoff "${dev}"
      rm -f "${RUN}/swap"
    else
      if [ ! -f "${RUN}/${dir//\//_}" ]; then
        printf '%s is already unmounted from %s\n' "${fs}" "${dir}"
        continue
      fi
      dev="$(cat ${RUN}/${dir//\//_})"
      printf 'Stopping %s (mounted on %s)\n' "${dev}" "${dir}"
      umount "${dir}" || :   # Ignore unmount failures
      rm -f "${RUN}/${dir//\//_}"
    fi
    zramctl "${dev}" -r
  done
}


main() {
  if [ ${UID} -ne 0 ]; then
    printf 'Must be run as root\n' >&2
    return 1
  fi

  # Make sure the zram module is loaded
  modprobe zram

  if [ "${1:-}" = 'start' ]; then
    start
  elif [ "${1:-}" = 'stop' ]; then
    stop
  elif [ "${1:-}" = 'status' ]; then
    status
  else
    printf 'usage: %s [start|stop|status]\n' "${0}"
    return 1
  fi
}

main ${@}