#!/bin/bash

## Copyright (C) 2024 - 2025 ENCRYPTED SUPPORT LLC <adrelanos@whonix.org>
## See the file COPYING for copying conditions.

set -e
set -o pipefail
set -o nounset

calculate_swap_size() {
  local ram_in_mb=$1
  local disk_in_mb=$2
  local hibernation=$3
  local swap_size=0
  local sqrt_ram_gb
  local sqrt_ram_mb
  local disk_max_swap=0

  echo "Calculating swap size..."
  echo "RAM in MB: $ram_in_mb"
  echo "Disk size in MB: $disk_in_mb"
  echo "Hibernation considered: $hibernation"

  # Convert MB to GB for comparison
  local ram_in_gb=$((ram_in_mb / 1024))
  echo "RAM in GB (approximated for calculation): $ram_in_gb"

  if (( ram_in_gb <= 2 )); then
    if [[ "$hibernation" == "yes" ]]; then
      #swap_size=$((3 * ram_in_mb))
      #echo "RAM <= 2GB with hibernation: Swap size is 3 times the RAM"
      echo "RAM <= 2GB with hibernation: Swap size is at least 6 GB"
      swap_size=6144
    else
      #swap_size=$((2 * ram_in_mb))
      #echo "RAM <= 2GB without hibernation: Swap size is 2 times the RAM"
      echo "RAM <= 2GB without hibernation: Swap size is at least 6 GB"
      swap_size=6144
    fi
  elif (( ram_in_gb > 2 && ram_in_gb <= 8 )); then
    if [[ "$hibernation" == "yes" ]]; then
      swap_size=$((2 * ram_in_mb))
      echo "RAM > 2GB and <= 8GB with hibernation: Swap size is 2 times the RAM"
    else
      swap_size=$ram_in_mb
      echo "RAM > 2GB and <= 8GB without hibernation: Swap size equals RAM"
    fi
  elif (( ram_in_gb > 8 && ram_in_gb <= 64 )); then
    if [[ "$hibernation" == "yes" ]]; then
      swap_size=$((ram_in_mb + ram_in_mb / 2)) # 1.5 times the RAM
      echo "RAM > 8GB and <= 64GB with hibernation: Swap size is 1.5 times the RAM"
    else
      swap_size=$(( ram_in_mb / 5 ))
      echo "RAM > 8GB and <= 64GB without Hibernation: Swap size is 20% of the RAM"
    fi
  else
    if [[ "$hibernation" == "yes" ]]; then
      # Calculate the square root of RAM size in GB
      sqrt_ram_gb=$(echo "scale=0; sqrt($ram_in_gb)" | bc)
      # Convert the square root from GB to MB by multiplying by 1024
      sqrt_ram_mb=$((sqrt_ram_gb * 1024))
      # Add the square root in MB to the RAM size in MB for the swap size
      swap_size=$((ram_in_mb + sqrt_ram_mb))
      echo "RAM > 64GB with Hibernation: Swap size is size of the RAM ($ram_in_mb MB) + square root of RAM ($sqrt_ram_mb MB) ($sqrt_ram_gb GB)"
    else
      swap_size=$(( ram_in_mb / 5 ))
      echo "RAM > 64GB without Hibernation: Swap size is 20% of the RAM"
    fi
  fi

  disk_max_swap=$(( disk_in_mb / 10 ))
  if (( disk_max_swap < swap_size )); then
    swap_size=$disk_max_swap
    echo "Capping swap size at $swap_size (10% of total disk space)"
  fi

  echo "Calculated Swap Size in MB: $swap_size"

  CALCULATED_SWAP_SIZE="$swap_size"
}

test_calculate_failed() {
  echo "ERROR: test failed number: $1" >&2
  exit 1
}

test_calculate_swap_size() {
  # The first ten tests all have a 4 TB disk specified to avoid running into caps
  # Test 1: RAM <= 2GB without hibernation
  calculate_swap_size 2048 4194304 no
  echo "Test 1: 2 GB RAM without hibernation: Expected 6144, Got $CALCULATED_SWAP_SIZE"
  (( CALCULATED_SWAP_SIZE == 6144 )) || test_calculate_failed 1

  # Test 2: RAM <= 2GB with hibernation
  calculate_swap_size 2048 4194304 yes
  echo "Test 2: 2 GB RAM with hibernation: Expected 6144, Got $CALCULATED_SWAP_SIZE"
  (( CALCULATED_SWAP_SIZE == 6144 )) || test_calculate_failed 2

  # Test 3: RAM > 2GB and <= 8GB without hibernation
  calculate_swap_size 4096 4194304 no
  echo "Test 3: 4 GB RAM without hibernation: Expected 4096, Got $CALCULATED_SWAP_SIZE"
  (( CALCULATED_SWAP_SIZE == 4096 )) || test_calculate_failed 3

  # Test 4: RAM > 2GB and <= 8GB with hibernation
  calculate_swap_size 4096 4194304 yes
  echo "Test 4: 4 GB RAM with hibernation: Expected 8192, Got $CALCULATED_SWAP_SIZE"
  (( CALCULATED_SWAP_SIZE == 8192 )) || test_calculate_failed 4

  # Test 5: RAM > 8GB and <= 64GB without hibernation
  calculate_swap_size 16384 4194304 no
  echo "Test 5: 16 GB RAM without hibernation: Expected 3276, Got $CALCULATED_SWAP_SIZE"
  (( CALCULATED_SWAP_SIZE == 3276 )) || test_calculate_failed 5

  # Test 6: RAM > 8GB and <= 64GB with hibernation
  calculate_swap_size 16384 4194304 yes
  echo "Test 6: 16 GB RAM with hibernation: Expected 24576, Got $CALCULATED_SWAP_SIZE"
  (( CALCULATED_SWAP_SIZE == 24576 )) || test_calculate_failed 6

  # Test 7: RAM > 64GB without hibernation
  calculate_swap_size 65536 4194304 no
  echo "Test 7: 64 GB RAM without hibernation: Expected 13107, Got $CALCULATED_SWAP_SIZE"
  (( CALCULATED_SWAP_SIZE == 13107 )) || test_calculate_failed 7

  # Test 8: RAM > 64GB with hibernation
  calculate_swap_size 65536 4194304 yes
  echo "Test 8: 64 GB RAM with hibernation: Expected 98304, Got $CALCULATED_SWAP_SIZE"
  (( CALCULATED_SWAP_SIZE == 98304 )) || test_calculate_failed 8

  # Test for 128GB RAM without hibernation
  calculate_swap_size 131072 4194304 no
  echo "Test 9: 128GB RAM without hibernation: Expected 26214, Got $CALCULATED_SWAP_SIZE"
  (( CALCULATED_SWAP_SIZE == 26214 )) || test_calculate_failed 9

  # Test for 128GB RAM with hibernation
  calculate_swap_size 131072 4194304 yes
  echo "Test 10: 128GB RAM with hibernation: Expected 142336, Got $CALCULATED_SWAP_SIZE"
  (( CALCULATED_SWAP_SIZE == 142336 )) || test_calculate_failed 10

  # Test for 16GB RAM with hibernation and a small disk
  calculate_swap_size 16384 65536 yes
  echo "Test 11: 16GB RAM with hibernation but only 64 GB disk: Expected 6553, Got $CALCULATED_SWAP_SIZE"
  (( CALCULATED_SWAP_SIZE == 6553 )) || test_calculate_failed 11
}

## For testing:
## Uncomment the following line:
#test_calculate_swap_size

if [[ $# -ne 3 ]] || [[ $# -ge 4 ]] ; then
    echo "Usage: $0 [RAM in MB] [Disk size in MB] [Hibernation (yes/no)]" >&2
    exit 1
fi

RAM_IN_MB="$1"
DISK_IN_MB="$2"
HIBERNATION="$3"

if [[ ! $RAM_IN_MB =~ ^[0-9]+$ ]] || [[ ! $DISK_IN_MB =~ ^[0-9]+$ ]] || { [[ $HIBERNATION != "yes" ]] && [[ $HIBERNATION != "no" ]]; }; then
  echo "Invalid input. Please provide the RAM and disk size in MB and specify hibernation as 'yes' or 'no'." >&2
  exit 1
fi

calculate_swap_size "$RAM_IN_MB" "$DISK_IN_MB" "$HIBERNATION"
