#!/usr/bin/env bash
# Copyright (C) 2024-2025 Vasiliy Stelmachenok for CachyOS
#
# 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 2 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, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
file_path="/sys/kernel/mm/ksm/run"

die() {
    echo "$@" >&2
    exit 1
}

help() {
    die "$(cat << EOF
Usage: $0 {--enable|-e|--disable|-d|--status|-s|--pause|-p|--help|-h}
  --enable,  -e :  Enable KSM
  --disable, -d :  Disable KSM
  --pause,   -p :  Temporarily disable KSM without unmerging pages
  --status,  -s :  Print the current KSM status
  --help,    -h :  Print this message
EOF
)"
}

if [ "$#" -eq 0 ]; then
    help
fi

while [[ "$#" -gt 0 ]]; do
    case "$1" in
        --pause | -p)
            new_state="0"
            shift
            ;;
        --enable | -e)
            new_state="1"
            shift
            ;;
        --disable | -d)
            new_state="2"
            shift
            ;;
        --help | -h)
            help
            ;;
        --status | -s)
            status="Unknown KSM state"

            case "$(cat "$file_path")" in
                0) status="KSM disabled, already merged tables kept" ;;
                1) status="KSM enabled and working" ;;
                2) status="KSM disabled, already merged tables unmerged" ;;
                *) ;;
            esac

            echo "$status"
            exit 0
            ;;
        *)
            die "Error: Invalid argument. Use --enable, --disable, --pause, --status or --help."
            ;;
    esac
done

if [ "$(id -u)" -ne 0 ]; then
    die "Error: This script must be run with root privileges."
fi

if [ -z "$new_state" ]; then
    die "Error: Missing action. Use --enable, --disable or --pause" >&2
fi

if ((new_state == 0)); then
    echo "Pausing Kernel Samepage Merging..."
    echo "0" > "$file_path"
elif ((new_state == 1)); then
    echo "Enabling Kernel Samepage Merging..."
    echo "1" > "$file_path"
else
    echo "Disabling Kernel Samepage Merging..."
    echo "2" > "$file_path"
fi