#!/bin/sh
# This script uses 'local' variables and expects a shell that supports them (e.g., dash)
# shellcheck disable=SC3043  # Ignore 'local' is undefined warnings

set -e

: "${CONFIG_FILE:=/etc/rpi/swap.conf}"
: "${SWAP_FILE_SIZE_CALC_BIN:=/usr/lib/rpi-swap/bin/rpi-desired-swap-size}"
: "${ZRAM_SIZE_CALC_BIN:=/usr/lib/rpi-swap/bin/rpi-desired-zram-size}"

# Function to create a service that removes the swap file
create_remove_swap_file_service() {
  local swap_file="$1"

  cat <<EOF > /run/systemd/generator/rpi-remove-swap-file.service
# Automatically generated by rpi-swap-generator

[Unit]
SourcePath=${CONFIG_FILE}
DefaultDependencies=no
After=systemd-remount-fs.service
Requires=systemd-remount-fs.service

[Service]
Type=oneshot
ExecStart=rm -f ${swap_file}
EOF
  mkdir -p /run/systemd/generator/local-fs.target.wants
  ln -s "../rpi-remove-swap-file.service" /run/systemd/generator/local-fs.target.wants/rpi-remove-swap-file.service
}

# Function to create a swap unit file
create_swap_unit() {
  local what_device="$1"
  local unit_depends="$2"

  # Derive the unit filename from the device path
  local unit_filename
  unit_filename="$(systemd-escape --path --suffix=swap "${what_device}")"
  local unit_path="/run/systemd/generator/${unit_filename}"

  cat <<EOF > "${unit_path}"
# Automatically generated by rpi-swap-generator

[Unit]
SourcePath=${CONFIG_FILE}
Requires=${unit_depends}
After=${unit_depends}

[Swap]
What=${what_device}
EOF

  # Create directory for the symlink if it doesn't exist
  mkdir -p /run/systemd/generator/swap.target.wants

  # Create symlink to enable the unit
  ln -s "../${unit_filename}" /run/systemd/generator/swap.target.wants/"${unit_filename}"
}

# Function to create systemd drop-in configuration files
create_systemd_drop_in() {
  local base_path="$1"         # e.g. 'generator' or 'zram-generator.conf.d'
  local target_name="$2"       # e.g. 'systemd-zram-setup@zram0.service' or empty for base config
  local config_name="$3"       # e.g. 'bindings.conf' or '80-rpi-enable.zram.conf'
  local content="$4"           # The content to write to the file

  local dir_path
  local file_path

  if [ -n "${target_name}" ]; then
    # Creating a drop-in for a specific service
    dir_path="/run/systemd/${base_path}/${target_name}.d"
    file_path="${dir_path}/${config_name}"
  else
    # Creating a conf file directly in the base_path
    dir_path="/run/systemd/${base_path}"
    file_path="${dir_path}/${config_name}"
  fi

  # Create directory if it doesn't exist
  mkdir -p "${dir_path}"

  # Create the configuration file
  cat <<EOF > "${file_path}"
# Automatically generated by rpi-swap-generator

${content}
EOF
}

# Function to create writeback timer and service units
create_writeback_units() {
  local initial_delay="$1"
  local periodic_interval="$2"

  # Create the service unit
  cat <<EOF > /run/systemd/generator/rpi-zram-writeback.service
# Automatically generated by rpi-swap-generator

[Unit]
SourcePath=${CONFIG_FILE}
Description=zram writeback service
After=dev-zram0.swap

[Service]
Type=oneshot
ExecStart=/bin/sh -c 'echo idle > /sys/block/zram0/writeback'
EOF

  # Create the timer unit
  cat <<EOF > /run/systemd/generator/rpi-zram-writeback.timer
# Automatically generated by rpi-swap-generator

[Unit]
SourcePath=${CONFIG_FILE}
Description=zram writeback timer

[Timer]
OnBootSec=${initial_delay}
OnUnitActiveSec=${periodic_interval}

[Install]
WantedBy=timers.target
EOF

  # Create symlink to enable the timer
  mkdir -p /run/systemd/generator/timers.target.wants
  ln -s "../rpi-zram-writeback.timer" /run/systemd/generator/timers.target.wants/rpi-zram-writeback.timer
}

# Function to set up zram mechanisms (both zram and zram+file)
setup_zram_mechanism() {
  local mechanism="$1"  # Either "zram" or "zram+file"
  local zram_config="[zram0]
host-memory-limit=none
fs-type=swap"

  # Add zram size if calculated
  if [ -n "${CONF_ZRAM_SIZE}" ] && [ "${CONF_ZRAM_SIZE}" -gt 0 ] 2>/dev/null; then
    # systemd-zram-generator expects the value in MiB as a mathematical expression
    zram_config="${zram_config}
zram-size=${CONF_ZRAM_SIZE}"
  fi
  local bindings="[Unit]
BindsTo=dev-%i.swap"

  # For zram+file, set up backing file and add to config
  if [ "${mechanism}" = "zram+file" ]; then
    BACKING_FILE="$( \
      systemd-escape \
        --path \
        "${CONF_SWAPFILE}" \
    )"
    LOOP_DEVICE="/dev/disk/by-backingfile/${BACKING_FILE}"

    # Add writeback device to zram config
    zram_config="${zram_config}
writeback-device=${LOOP_DEVICE}"

    # Add loop device to bindings
    bindings="${bindings} rpi-setup-loop@${BACKING_FILE}.service"

    # Create a drop-in for the loop service to ensure the swap file is resized
    create_systemd_drop_in \
      "generator" \
      "rpi-setup-loop@${BACKING_FILE}.service" \
      "resize-swap.conf" \
      "[Unit]
Requires=rpi-resize-swap-file.service
After=rpi-resize-swap-file.service"
  fi

  # Create zram configuration
  create_systemd_drop_in \
    "zram-generator.conf.d" \
    "" \
    "20-rpi-swap-zram0-ctrl.conf" \
    "${zram_config}"

  # Create a swap unit
  create_swap_unit "/dev/zram0" "systemd-zram-setup@zram0.service"

  # Create the bindings
  create_systemd_drop_in \
    "generator" \
    "systemd-zram-setup@zram0.service" \
    "bindings.conf" \
    "${bindings}"

  # Create the ordering for zram+file mechanism
  if [ "${mechanism}" = "zram+file" ]; then
    create_systemd_drop_in \
      "generator" \
      "systemd-zram-setup@zram0.service" \
      "ordering.conf" \
      "[Unit]
After=rpi-setup-loop@${BACKING_FILE}.service"
  fi

  # For zram+file, we must mark the pages as idle before they're first used
  if [ "${mechanism}" = "zram+file" ]; then
    create_systemd_drop_in \
      "generator" \
      "systemd-zram-setup@zram0.service" \
      "mark-pages-idle.conf" \
      "[Service]
ExecStartPost=/bin/sh -c 'echo all > /sys/block/zram0/idle'"
  fi
}

: "${CONF_SWAPFILE:=/var/swap}"
eval "$( \
  rpi-systemd-config rpi/swap.conf \
    CONF_MECHANISM     Main::Mechanism \
    CONF_SWAPFILE      File::Path \
    CONF_SWAPFACTOR    File::RamMultiplier \
    CONF_MAXDISK_PCT   File::MaxDiskPercent \
    CONF_MAXSWAP       File::MaxSizeMiB \
    CONF_SWAPSIZE      File::FixedSizeMiB \
    CONF_ZRAM_FACTOR   Zram::RamMultiplier \
    CONF_ZRAM_MAXSIZE  Zram::MaxSizeMiB \
    CONF_ZRAM_SIZE     Zram::FixedSizeMiB \
    CONF_WRITEBACK_INITIATOR      Zram::WritebackTrigger \
    CONF_INITIAL_WRITEBACK_DELAY  Zram::WritebackInitialDelay \
    CONF_PERIODIC_WRITEBACK_INTERVAL Zram::WritebackPeriodicInterval \
)"

# Set default mechanism to "auto" if not specified in config
: "${CONF_MECHANISM:=auto}"

# Set default writeback settings if not specified in config
: "${CONF_WRITEBACK_INITIATOR:=auto}"
: "${CONF_INITIAL_WRITEBACK_DELAY:=180min}"
: "${CONF_PERIODIC_WRITEBACK_INTERVAL:=24h}"

# Sanitize mechanism - ensure it's one of the allowed values
case "${CONF_MECHANISM}" in
  swapfile|zram|zram+file|auto|none)
    # Valid mechanism, do nothing
    ;;
  *)
    # Invalid mechanism, default to auto
    CONF_MECHANISM="auto"
    ;;
esac

# Sanitize writeback initiator - ensure it's one of the allowed values
case "${CONF_WRITEBACK_INITIATOR}" in
  auto|manual|timer)
    # Valid initiator, do nothing
    ;;
  *)
    # Invalid initiator, default to auto
    CONF_WRITEBACK_INITIATOR="auto"
    ;;
esac

# Handle auto mechanism
if [ "${CONF_MECHANISM}" = "auto" ]; then
  CONF_MECHANISM="zram+file"
fi

# Handle auto writeback initiator
if [ "${CONF_WRITEBACK_INITIATOR}" = "auto" ]; then
  CONF_WRITEBACK_INITIATOR="timer"
fi

# For 'zram' and 'none' mechanisms, we don't need a swap file at all
if [ "${CONF_MECHANISM}" = "zram" ] || [ "${CONF_MECHANISM}" = "none" ]; then
  # Skip calculation and set swap size to 0
  CONF_SWAPSIZE="0"
else
  # Calculate swap size for other mechanisms
  export \
    CONF_SWAPFILE \
    CONF_SWAPFACTOR \
    CONF_MAXDISK_PCT \
    CONF_MAXSWAP \
    CONF_SWAPSIZE
  CONF_SWAPSIZE="$(${SWAP_FILE_SIZE_CALC_BIN})"
fi

# Calculate zram size for 'zram' and 'zram+file' mechanisms
if [ "${CONF_MECHANISM}" = "zram" ] || [ "${CONF_MECHANISM}" = "zram+file" ]; then
  export \
    CONF_ZRAM_FACTOR \
    CONF_ZRAM_MAXSIZE \
    CONF_ZRAM_SIZE
  CONF_ZRAM_SIZE="$(${ZRAM_SIZE_CALC_BIN})"
else
  CONF_ZRAM_SIZE="0"
fi

CONF_SWAPFILE="$(realpath "${CONF_SWAPFILE}")"

if [ "${CONF_MECHANISM}" = "zram" ]; then
  setup_zram_mechanism "zram"
  if [ -f "${CONF_SWAPFILE}" ]; then
    create_remove_swap_file_service "${CONF_SWAPFILE}"
  fi
elif [ "${CONF_MECHANISM}" = "zram+file" ]; then
  setup_zram_mechanism "zram+file"
  if [ "${CONF_WRITEBACK_INITIATOR}" = "timer" ]; then
    create_writeback_units "${CONF_INITIAL_WRITEBACK_DELAY}" "${CONF_PERIODIC_WRITEBACK_INTERVAL}"
  fi
elif [ "${CONF_MECHANISM}" = "none" ]; then
  # No swap setup at all, but clean up existing swap file if present
  if [ -f "${CONF_SWAPFILE}" ]; then
    create_remove_swap_file_service "${CONF_SWAPFILE}"
  fi
elif [ "${CONF_SWAPSIZE}" -ne 0 ]; then # Assumed 'swapfile' case
  create_swap_unit "${CONF_SWAPFILE}" "rpi-resize-swap-file.service"
elif [ -f "${CONF_SWAPFILE}" ]; then
  create_remove_swap_file_service "${CONF_SWAPFILE}"
fi
