#!/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

# Set generator directory variables from command line arguments
GEN_NORMAL_DIR="$1"
GEN_EARLY_DIR="$2"
GEN_LATE_DIR="$3"

if [ -z "${GEN_NORMAL_DIR}" ]; then
  echo "Error: Missing generator directory argument" >&2
  exit 1
fi

: "${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 symlinks safely - only fails if existing symlink points elsewhere
create_symlink() {
  local target="$1"
  local link_path="$2"

  if [ -L "${link_path}" ]; then
    # Symlink exists, check if it points to the right target
    local existing_target
    existing_target="$(readlink "${link_path}")"
    if [ "${existing_target}" != "${target}" ]; then
      echo "Error: Symlink ${link_path} exists but points to ${existing_target}, not ${target}" >&2
      return 1
    fi
    # Symlink already points to correct target, nothing to do
    return 0
  elif [ -e "${link_path}" ]; then
    # Path exists but is not a symlink
    echo "Error: ${link_path} exists but is not a symlink" >&2
    return 1
  else
    # Create the symlink
    ln -s "${target}" "${link_path}"
  fi
}

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

  # Use systemd-escape to create a safe instance name from the file path
  local escaped_path
  escaped_path="$(systemd-escape --path "${swap_file}")"
  local unit_name="rpi-remove-swap-file@${escaped_path}.service"

  # Create symlink to enable the templated service
  mkdir -p "${GEN_NORMAL_DIR}/local-fs.target.wants"
  create_symlink "../${unit_name}" "${GEN_NORMAL_DIR}/local-fs.target.wants/${unit_name}"
}

# Function to create a swap unit file
create_swap_unit() {
  local what_device="$1"
  local unit_depends="$2"
  local conflict_file_path="$3"  # Optional: file path to conflict with (unescaped)

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

  # Build conflicts line for file removal service
  local conflicts_line=""
  if [ -n "${conflict_file_path}" ]; then
    local escaped_conflict_file
    escaped_conflict_file="$(systemd-escape --path "${conflict_file_path}")"
    conflicts_line="Conflicts=rpi-remove-swap-file@${escaped_conflict_file}.service"
  fi

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

[Unit]
SourcePath=${CONFIG_FILE}
Description=rpi-swap managed swap device (${CONF_MECHANISM})
Requires=${unit_depends}
After=${unit_depends}
${conflicts_line}

[Swap]
What=${what_device}
Priority=100
EOF

  # Create directory for the symlink if it doesn't exist
  mkdir -p "${GEN_NORMAL_DIR}/swap.target.wants"

  # Create symlink to enable the unit
  create_symlink "../${unit_filename}" "${GEN_NORMAL_DIR}/swap.target.wants/${unit_filename}"
}

# Function to create systemd drop-in configuration files
create_systemd_drop_in() {
  local base_path="$1"         # Special values: 'generator', 'generator.early', 'generator.late'
                               # or standard systemd paths like '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
  local generator_dir

  # Map special base_path values to appropriate generator directories
  case "${base_path}" in
    "generator")
      generator_dir="${GEN_NORMAL_DIR}"
      ;;
    "generator.early")
      generator_dir="${GEN_EARLY_DIR}"
      ;;
    "generator.late")
      generator_dir="${GEN_LATE_DIR}"
      ;;
    *)
      # For any other value, use standard /run/systemd/ prefix
      generator_dir="/run/systemd/${base_path}"
      ;;
  esac

  if [ -n "${target_name}" ]; then
    # Creating a drop-in for a specific service
    dir_path="${generator_dir}/${target_name}.d"
    file_path="${dir_path}/${config_name}"
  else
    # Creating a conf file directly in the base_path
    dir_path="${generator_dir}"
    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 unit (service unit is static)
create_writeback_units() {
  local initial_delay="$1"
  local periodic_interval="$2"


  # Create the timer unit (service unit is shipped as static file)
  cat <<EOF > "${GEN_NORMAL_DIR}/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 "${GEN_NORMAL_DIR}/timers.target.wants"
  create_symlink "../rpi-zram-writeback.timer" "${GEN_NORMAL_DIR}/timers.target.wants/rpi-zram-writeback.timer"
}

# Function to create sysctl configuration for zram optimizations
create_zram_sysctl_config() {
  local sysctl_dir="/run/sysctl.d"
  local sysctl_file="${sysctl_dir}/70-rpi-swap.conf"

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

  # Create the sysctl configuration file (idempotent)
  cat <<EOF > "${sysctl_file}"
# Automatically generated by rpi-swap-generator

# Disable readahead for zram swap devices
# Reduces latency with minimal throughput impact for compressed RAM swap
vm.page-cluster=0
EOF
}

# Function to set up zram mechanisms (both zram and zram+file)
setup_zram_mechanism() {
  local mechanism="$1"  # Either "zram" or "zram+file"

  # Create sysctl optimizations for zram
  create_zram_sysctl_config

  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}" \
    )"
    local 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"

    # Create a drop-in to avoid dependency cycles with swap.target
    create_systemd_drop_in \
      "generator" \
      "rpi-setup-loop@${BACKING_FILE}.service" \
      "no-dependency-cycle.conf" \
      "[Unit]
DefaultDependencies=no
Requires=systemd-remount-fs.service
After=systemd-remount-fs.service"

    # Create a drop-in to speed up symlink creation for swap use case
    create_systemd_drop_in \
      "generator" \
      "rpi-setup-loop@${BACKING_FILE}.service" \
      "fast-symlink.conf" \
      "[Service]
ExecStart=/bin/sh -c 'mkdir -p /dev/disk/by-backingfile && LOOP_DEV=\$(losetup --associated \"/%I\" | cut -d: -f1) && ln -s \"\$LOOP_DEV\" /dev/disk/by-backingfile/%i || true'"
  fi

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

  # Create a swap unit with backing file conflict for zram+file
  local backing_file=""
  if [ "${mechanism}" = "zram+file" ]; then
    backing_file="${CONF_SWAPFILE}"
  fi
  create_swap_unit "/dev/zram0" "systemd-zram-setup@zram0.service" "${backing_file}"

  # 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" "${CONF_SWAPFILE}"
elif [ -f "${CONF_SWAPFILE}" ]; then
  create_remove_swap_file_service "${CONF_SWAPFILE}"
fi
