#!/bin/bash
# Lackadaisical squashfs tools
# Allows you to create a modifiable squashfs-backed image for a folder
BIN_DIR=$(dirname "$(readlink -f "$0")")
DIR=$1

if [[ -z "$DIR" ]]; then
  echo "Usage: $0 <directory>"
  exit 1
fi

test -d "$DIR"
if [[ $? -ne 0 ]]; then
  echo Directory "$DIR" does not exist!
  exit -1
fi

echo "You need support for FUSE, SquashFS, and OverlayFS."
echo "Checking your system for these requirements..."
for fs in fuse squashfs overlay; do
  if ! grep -q "$fs" /proc/filesystems; then
    echo "Loading $fs kernel module..."
    if ! sudo modprobe "$fs"; then
      echo "Error: $fs is not supported and could not be loaded."
      exit 1
    fi
  fi
done

DIRSIZE=$(du -s "$DIR")

DIR=$(readlink -f "$DIR")
DIR_SHORT=$(basename "$DIR")
echo Compressing "$DIR"

# Make basic dirs
mkdir -p "${DIR}"
mkdir -p "${DIR_SHORT}"

OVERLAY_ROOT=$(readlink -f "${DIR}/..")/.squashfs/${DIR_SHORT}
OVERLAY_UPPER=${OVERLAY_ROOT}/upper
OVERLAY_LOWER=${OVERLAY_ROOT}/lower
OVERLAY_WORK=${OVERLAY_ROOT}/work
OVERLAY_TARG=$DIR
RECREATE=false

echo "Overlay information"
echo "========================================================================="
echo "> DIR = \"$DIR\""
echo "> DIR_SHORT = \"$DIR_SHORT\""
echo "> OVERLAY_ROOT = \"$OVERLAY_ROOT\""
echo "> OVERLAY_UPPER = \"$OVERLAY_UPPER\""
echo "> OVERLAY_LOWER= \"$OVERLAY_LOWER\""
echo "> OVERLAY_WORK = \"$OVERLAY_WORK\""
echo "> OVERLAY_TARG = \"$OVERLAY_TARG\""
echo "========================================================================="

mkdir -p "$OVERLAY_ROOT"

# Check for existing image
if test -f "${OVERLAY_ROOT}.img"; then
  echo "We already have an existing image, updating it..."
  mount-squash-image "$DIR"
  DIRSIZE=$(du -s "$DIR")
  RECREATE=true
fi

sudo mksquashfs "$DIR" "${OVERLAY_ROOT}.img.1" -noappend -comp xz

if [ $RECREATE == true ]; then
  echo "Copy created, now unmount and swap everything."
  umount-squash-image "$DIR"
  sudo rm -rf "$OVERLAY_UPPER"
  sudo rm -rf "$OVERLAY_LOWER"
  sudo rm -rf "$OVERLAY_WORK"
  rm -rf "${OVERLAY_ROOT}.img"
fi

mv "${OVERLAY_ROOT}.img.1" "${OVERLAY_ROOT}.img"

# Make the dirs for the overlay
mkdir -p "$OVERLAY_UPPER"
mkdir -p "$OVERLAY_LOWER"
mkdir -p "$OVERLAY_WORK"
mkdir -p "$OVERLAY_TARG"

# Reset dir
sudo rm -rf "$DIR"
mkdir -p "$DIR"
touch "${DIR}/.needs_mount"

echo "========================================================================="
echo "Storage data:"
echo "Uncompressed: $DIRSIZE"
echo "   Compressed: $(du -s ${OVERLAY_ROOT}.img)"
echo "========================================================================="

# Service content embedded
SERVICE_CONTENT="[Unit]
Description=SquashFS Mount for %I
After=local-fs.target

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=${BIN_DIR}/mount-squash-image %I
ExecStop=${BIN_DIR}/umount-squash-image %I

[Install]
WantedBy=multi-user.target"

echo "$SERVICE_CONTENT" | sudo tee /etc/systemd/system/squash-mount@.service > /dev/null
sudo systemctl daemon-reload

ESC_PATH=$(systemd-escape -p "$DIR")
SERVICE_NAME="squash-mount@-$ESC_PATH"

read -p "Do you want to auto-enable the mount service ($SERVICE_NAME)? [y/N] " yn
case $yn in
    [Yy]* ) 
        echo "Enabling and starting service..."
        sudo systemctl enable --now "$SERVICE_NAME"
        sudo systemctl stop "$SERVICE_NAME"
        sudo systemctl start "$SERVICE_NAME"
        ;;
    * ) 
        echo "To enable the service manually, run: sudo systemctl enable --now $SERVICE_NAME"
        ;;
esac

echo ""
echo "Your SquashFS-backed folder is ready for use."
echo "To mount it manually, either cd into it or use mount-squash-image."
echo ""
echo "Should you wish to update the contents of the image with your changes made"
echo "inside of the folder, simply run make-squash-image again on the same"
echo "folder to update the image."
echo
echo "To disable auto-mounting upon cd, please remove the"
echo ".needs_mount file to prevent that."
echo