#!/bin/bash

showusage () {
 echo "Usage: zramsetup <option>"
 echo "options:"
 echo "-i <max devices> - perform zram module init with selected maximum number of devices"
 echo "-I - perform zram module init with 10 devices"
 echo "-A <algo> - set compression algorithm. defaults to lz4. pass this option prior to -c option"
 echo "-S <count> - set compression streams. defaults to 1. pass this option prior to -c option"
 echo "-c <size> - prepare first available device for use with selected size and return device name on success"
 echo "-d <device name or number> - free and deinitialize selected device (by full device name or number)"
 echo "-D - perform zram module deinit"
 echo "-h - show help"
}

check_errors () {
  local status="$?"
  local msg="$1"
  if [ "$status" != "0" ]; then
    if [ "zzz$msg" = "zzz" ]; then
      echo "error detected in last operation"
    else
      echo "$msg"
    fi
    exit 1
  fi
}

is_integer () {
 printf "%d" $1 > /dev/null 2>&1
 return $?
}

waittime="3"

wait_for_zram () {
 local try="0"
 local result="fail"
 while test $try -lt $waittime ; do
   if [ -b "/dev/zram$1" ] ; then
     result="ok"
     break
   fi
   try=`expr $try + 1`
   sleep 1
 done
 echo $result
}

module_init () {
  local mdevs="$1"
  if [ "zzz$mdevs" != "zzz" ]; then
    if ! is_integer "$mdevs"; then
     echo "max devices parameter is not integer"
     exit 1
    fi
  else
    mdevs="10"
  fi
  
  local check=`lsmod | grep zram`
  if [ "zzz$check" = "zzz" ]; then
    modprobe zram num_devices=$mdevs 2> /dev/null
    check_errors "modprobe zram num_devices=$mdevs failed"
    
    local limit=`expr $mdevs - 1`
    for i in `seq 0 $limit`
    do
      check=`wait_for_zram $i`
      check_errors "failed to check $i zram device"
      if [ "$check" != "ok" ]; then
        echo "zram device $i detection failed"
        exit 1
      fi
    done
  fi
}

check_device () {
  local result="fail"
  local check=`cat /sys/block/zram$1/disksize 2> /dev/null`
  if [ "zzz$check" = "zzz0" ]; then
    result="ok"
  else
    if [ "zzz$check" = "zzz" ]; then
      result="null"
    fi
  fi
  echo $result
}

find_device () {
  local mdev="0"
  local ch_result="zzz"
  local result="fail"
  while true; do
    ch_result=`check_device $mdev`
    
    if [ "$ch_result" = "ok" ]; then
      result="$mdev"
      break
    fi
    
    if [ "$ch_result" = "null" ]; then
      break
    fi
    
    mdev=`expr $mdev + 1`
  done
  echo "$result"
}

init_device () {
  local mdev="$1"
  local size="$2"
  local streams="$3"
  local compr="$4"

  local check=`check_device $mdev`
  if [ "$check" != "ok" ]; then
    echo "error while performing init of zram$mdev device"
    exit 1
  fi
  if ! is_integer "$size"; then
    echo "device size parameter is not integer"
    exit 1
  fi
  if ! is_integer "$streams"; then
    echo "streams parameter is not integer"
    exit 1
  fi
  if [ "zzz$compr" = "zzz" ]; then
    echo "comp_algorithm parameter is incorrect"
    exit 1
  fi
  
  if [ -f "/sys/block/zram$mdev/max_comp_streams" ]; then
    echo "$streams" > /sys/block/zram$mdev/max_comp_streams
    check_errors "failed to set compression stream count for zram$mdev device"
  fi
  
  if [ -f "/sys/block/zram$mdev/comp_algorithm" ]; then
    echo "$compr" > /sys/block/zram$mdev/comp_algorithm
    check_errors "failed to set compression algorithm for zram$mdev device"
  fi
  
  echo "$size" > /sys/block/zram$mdev/disksize
  check_errors "failed to init size for zram$mdev device"
}

deinit_device () {
  local mdev="$1"
  if is_integer "$mdev"; then
    mdev="zram$mdev"
  else
    mdev=`basename $mdev`
  fi
  echo 1 2> /dev/null 1> /sys/block/$mdev/reset
  check_errors "failed to perform deinit on $mdev device"
}

module_deinit () {  
  local check=`lsmod | grep zram`
  if [ "zzz$check" != "zzz" ]; then
    modprobe -r zram 2> /dev/null
    check_errors "modprobe -r zram failed"
  fi
}

c_streams="1"
c_alg="lz4"

parseopts () {
 local optname
 while getopts ":i:Ic:d:DhS:A:" optname
 do
  case "$optname" in
   "i")
     module_init $OPTARG
     exit 0
   ;;
   
   "I")
     module_init
     exit 0
   ;;
   
   "c")
    module_init
    local devnum=`find_device`
    if [ "zzz$devnum" = "zzzfail" ]; then
      echo "no more spare devices left!"
      exit 1
    fi
    init_device $devnum $OPTARG "$c_streams" "$c_alg"
    echo "/dev/zram$devnum"
    exit 0
   ;;
   
   "d")
    module_init
    deinit_device $OPTARG
    exit 0
   ;;
   
   "D")
    module_deinit
    exit 0
   ;;
   
   "h")
    showusage
    exit 1
   ;;

   "S")
     c_streams="$OPTARG"
   ;;

   "A")
     c_alg="$OPTARG"
   ;;
   
   "?")
    echo "Unknown option $OPTARG"
    showusage
    exit 1
   ;;
   
   *)
    echo "Error while processing options"
    showusage
    exit 1
   ;;
  esac
 done
}

parseopts "$@"
showusage
exit 1
