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

if [[ ! -d "$DIR" ]]; then
  echo "Error: Directory \"$DIR\" does not exist!"
  exit 1
fi

echo "Checking system requirements (FUSE, SquashFS, OverlayFS)..."
for fs in fuse squashfs overlay; do
  if ! grep -q "$fs" /proc/filesystems; then
    echo "Attempting to load $fs module..."
    sudo modprobe "$fs" || { echo "Error: $fs is not supported."; exit 1; }
  fi
done

DIR=$(readlink -f "$DIR")
DIR_SHORT=$(basename "$DIR")
DIRSIZE=$(du -sh "$DIR" | cut -f1)

OVERLAY_ROOT="$(dirname "$DIR")/.squashfs/$DIR_SHORT"
OVERLAY_UPPER="$OVERLAY_ROOT/upper"
OVERLAY_LOWER="$OVERLAY_ROOT/lower"
OVERLAY_WORK="$OVERLAY_ROOT/work"
OVERLAY_TARG="$DIR"
RECREATE=false

mkdir -p "$OVERLAY_ROOT"

if [[ -f "${OVERLAY_ROOT}.img" ]]; then
  echo "Existing image found, updating..."
  "$BIN_DIR/mount-squash-image" "$DIR"
  DIRSIZE=$(du -sh "$DIR" | cut -f1)
  RECREATE=true
fi

echo "Compressing \"$DIR\"..."
sudo mksquashfs "$DIR" "${OVERLAY_ROOT}.img.1" -noappend -comp xz || exit 1

if [[ "$RECREATE" == "true" ]]; then
  echo "Cleaning up old layers..."
  "$BIN_DIR/umount-squash-image" "$DIR"
  # Safety check before destructive RM
  if [[ -n "$OVERLAY_ROOT" && "$OVERLAY_ROOT" != "/" ]]; then
    sudo rm -rf "$OVERLAY_UPPER" "$OVERLAY_LOWER" "$OVERLAY_WORK"
    rm -f "${OVERLAY_ROOT}.img"
  fi
fi

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

# Prepare overlay structure
mkdir -p "$OVERLAY_UPPER" "$OVERLAY_LOWER" "$OVERLAY_WORK" "$OVERLAY_TARG"

# Reset original directory to mountpoint
sudo rm -rf "$DIR"
mkdir -p "$DIR"
touch "$DIR/.needs_mount"

echo "-------------------------------------------------------------------------"
echo "Storage Stats:"
echo "  Original size: $DIRSIZE"
echo "  Compressed:    $(du -sh "${OVERLAY_ROOT}.img" | cut -f1)"
echo "-------------------------------------------------------------------------"

# SystemD Service Setup
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.service"

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