diff --git a/kernel/setup.sh b/kernel/setup.sh
index 68107e01..e688dbaf 100755
--- a/kernel/setup.sh
+++ b/kernel/setup.sh
@@ -1,50 +1,75 @@
 #!/bin/sh
-set -eux
+set -eu
 
 GKI_ROOT=$(pwd)
 
-echo "[+] GKI_ROOT: $GKI_ROOT"
+display_usage() {
+    echo "Usage: $0 [--cleanup | <commit-or-tag>]"
+    echo "  --cleanup:              Cleans up previous modifications made by the script."
+    echo "  <commit-or-tag>:        Sets up or updates the KernelSU to specified tag or commit."
+    echo "  -h, --help:             Displays this usage information."
+    echo "  (no args):              Sets up or updates the KernelSU environment to the latest tagged version."
+}
 
-if test -d "$GKI_ROOT/common/drivers"; then
-     DRIVER_DIR="$GKI_ROOT/common/drivers"
-elif test -d "$GKI_ROOT/drivers"; then
-     DRIVER_DIR="$GKI_ROOT/drivers"
+initialize_variables() {
+    if test -d "$GKI_ROOT/common/drivers"; then
+         DRIVER_DIR="$GKI_ROOT/common/drivers"
+    elif test -d "$GKI_ROOT/drivers"; then
+         DRIVER_DIR="$GKI_ROOT/drivers"
+    else
+         echo '[ERROR] "drivers/" directory not found.'
+         exit 127
+    fi
+
+    DRIVER_MAKEFILE=$DRIVER_DIR/Makefile
+    DRIVER_KCONFIG=$DRIVER_DIR/Kconfig
+}
+
+# Reverts modifications made by this script
+perform_cleanup() {
+    echo "[+] Cleaning up..."
+    [ -L "$DRIVER_DIR/kernelsu" ] && rm "$DRIVER_DIR/kernelsu" && echo "[-] Symlink removed."
+    grep -q "kernelsu" "$DRIVER_MAKEFILE" && sed -i '/kernelsu/d' "$DRIVER_MAKEFILE" && echo "[-] Makefile reverted."
+    grep -q "drivers/kernelsu/Kconfig" "$DRIVER_KCONFIG" && sed -i '/drivers\/kernelsu\/Kconfig/d' "$DRIVER_KCONFIG" && echo "[-] Kconfig reverted."
+    if [ -d "$GKI_ROOT/KernelSU" ]; then
+        rm -rf "$GKI_ROOT/KernelSU" && echo "[-] KernelSU directory deleted."
+    fi
+}
+
+# Sets up or update KernelSU environment
+setup_kernelsu() {
+    echo "[+] Setting up KernelSU..."
+    test -d "$GKI_ROOT/KernelSU" || git clone https://github.com/tiann/KernelSU && echo "[+] Repository cloned."
+    cd "$GKI_ROOT/KernelSU"
+    git stash && echo "[-] Stashed current changes."
+    if [ "$(git status | grep -Po 'v\d+(\.\d+)*' | head -n1)" ]; then
+        git checkout main && echo "[-] Switched to main branch."
+    fi
+    git pull && echo "[+] Repository updated."
+    if [ -z "${1-}" ]; then
+        git checkout "$(git describe --abbrev=0 --tags)" && echo "[-] Checked out latest tag."
+    else
+        git checkout "$1" && echo "[-] Checked out $1." || echo "[-] Checkout default branch"
+    fi
+    cd "$DRIVER_DIR"
+    ln -sf "$(realpath --relative-to="$DRIVER_DIR" "$GKI_ROOT/KernelSU/kernel")" "kernelsu" && echo "[+] Symlink created."
+
+    # Add entries in Makefile and Kconfig if not already existing
+    grep -q "kernelsu" "$DRIVER_MAKEFILE" || printf "\nobj-\$(CONFIG_KSU) += kernelsu/\n" >> "$DRIVER_MAKEFILE" && echo "[+] Modified Makefile."
+    grep -q "source \"drivers/kernelsu/Kconfig\"" "$DRIVER_KCONFIG" || sed -i "/endmenu/i\source \"drivers/kernelsu/Kconfig\"" "$DRIVER_KCONFIG" && echo "[+] Modified Kconfig."
+    echo '[+] Done.'
+}
+
+# Process command-line arguments
+if [ "$#" -eq 0 ]; then
+    initialize_variables
+    setup_kernelsu
+elif [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
+    display_usage
+elif [ "$1" = "--cleanup" ]; then
+    initialize_variables
+    perform_cleanup
 else
-     echo '[ERROR] "drivers/" directory is not found.'
-     echo '[+] You should modify this script by yourself.'
-     exit 127
+    initialize_variables
+    setup_kernelsu "$@"
 fi
-
-test -d "$GKI_ROOT/KernelSU" || git clone https://github.com/tiann/KernelSU
-cd "$GKI_ROOT/KernelSU"
-git stash
-if [ "$(git status | grep -Po 'v\d+(\.\d+)*' | head -n1)" ]; then
-     git checkout main
-fi
-git pull
-if [ -z "${1-}" ]; then
-    git checkout "$(git describe --abbrev=0 --tags)"
-else
-    git checkout "$1"
-fi
-cd "$GKI_ROOT"
-
-echo "[+] GKI_ROOT: $GKI_ROOT"
-echo "[+] Copy kernel su driver to $DRIVER_DIR"
-
-cd "$DRIVER_DIR"
-if test -d "$GKI_ROOT/common/drivers"; then
-     ln -sf "../../KernelSU/kernel" "kernelsu"
-elif test -d "$GKI_ROOT/drivers"; then
-     ln -sf "../KernelSU/kernel" "kernelsu"
-fi
-cd "$GKI_ROOT"
-
-echo '[+] Add kernel su driver to Makefile'
-
-DRIVER_MAKEFILE=$DRIVER_DIR/Makefile
-DRIVER_KCONFIG=$DRIVER_DIR/Kconfig
-grep -q "kernelsu" "$DRIVER_MAKEFILE" || printf "\nobj-\$(CONFIG_KSU) += kernelsu/\n" >> "$DRIVER_MAKEFILE"
-grep -q "kernelsu" "$DRIVER_KCONFIG" || sed -i "/endmenu/i\\source \"drivers/kernelsu/Kconfig\"" "$DRIVER_KCONFIG"
-
-echo '[+] Done.'