#!/bin/bash
# Lackadaisical squashfs tools - Squasher
# Combined utility for make, mount, umount, and destroy operations.

BIN_DIR=$(dirname "$(readlink -f "$0")")
COMMAND="$1"
DIR="$2"

usage()
{
  echo "Usage: $0 {make|mount|umount|destroy} <directory>"
  exit 1
}

if [[ -z "$COMMAND" || -z "$DIR" ]]; then
  usage
fi

DIR=$(readlink -f "$DIR")
DIR_SHORT=$(basename "$DIR")
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"

case "$COMMAND" in
  make)
    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

    DIRSIZE=$(du -sh "$DIR" | cut -f1)
    RECREATE=false
    mkdir -p "$OVERLAY_ROOT"

    if [[ -f "${OVERLAY_ROOT}.img" ]]; then
      echo "Existing image found, updating..."
      "$0" mount "$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..."
      "$0" umount "$DIR"
      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"
    mkdir -p "$OVERLAY_UPPER" "$OVERLAY_LOWER" "$OVERLAY_WORK" "$OVERLAY_TARG"

    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 "-------------------------------------------------------------------------"

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

    [Service]
    Type=oneshot
    RemainAfterExit=yes
    ExecStart=${BIN_DIR}/squasher mount %I
    ExecStop=${BIN_DIR}/squasher umount %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 "Enable auto-mount service ($SERVICE_NAME)? [y/N] " yn

    if [[ "$yn" =~ ^[Yy]$ ]]; then
      sudo systemctl enable --now "$SERVICE_NAME"
      sudo systemctl stop "$SERVICE_NAME"
      sudo systemctl start "$SERVICE_NAME"
    else
      echo "Manual mount command: sudo systemctl start $SERVICE_NAME"
    fi
    ;;

  mount)
    if [[ ! -f "${OVERLAY_ROOT}.img" ]]; then
      echo "Error: SquashFS image \"${OVERLAY_ROOT}.img\" not found." >&2
      exit 1
    fi

    "$0" umount "$DIR" 2>/dev/null
    mkdir -p "$OVERLAY_LOWER"
    sudo mount "${OVERLAY_ROOT}.img" "$OVERLAY_LOWER" -t squashfs -o loop

    if [[ $? -ne 0 ]]; then
      echo "Error: Failed to mount squashfs image." >&2
      exit 1
    fi

    sudo mount -t overlay none "$OVERLAY_TARG" \
      -o lowerdir="$OVERLAY_LOWER",upperdir="$OVERLAY_UPPER",workdir="$OVERLAY_WORK"

    if [[ $? -ne 0 ]]; then
      echo "Error: Failed to mount overlay." >&2
      sudo umount "$OVERLAY_LOWER" 2>/dev/null
      exit 1
    fi

    echo "SquashFS filesystem is mounted and ready."
    ;;

  umount)
    sudo umount -l -R "$OVERLAY_TARG" 2>/dev/null
    sudo umount -l -R "$OVERLAY_LOWER" 2>/dev/null

    if mountpoint -q "$OVERLAY_TARG" || mountpoint -q "$OVERLAY_LOWER"; then
      echo "Warning: Filesystem is still mounted. Check for open processes."
      exit 1
    fi

    echo "SquashFS filesystem has been unmounted."
    ;;

  destroy)
    if [[ ! -f "$DIR/.needs_mount" ]]; then
      if ! mountpoint -q "$DIR"; then
        echo "Error: $DIR is not a SquashFS directory."
        exit 1
      fi
    fi

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

    echo "Disabling service ($SERVICE_NAME)..."
    sudo systemctl stop "$SERVICE_NAME" 2>/dev/null
    sudo systemctl disable "$SERVICE_NAME" 2>/dev/null

    echo "Ensuring image is mounted to preserve data..."
    "$0" mount "$DIR" 1>/dev/null 2>/dev/null

    echo "Destroying image and restoring data..."
    TEMP_DIR=$(mktemp -d /tmp/squash-dest.XXXXXXXX)
    sudo rsync -aX "$DIR/" "$TEMP_DIR/" || { echo "Error: Failed to copy data."; exit 1; }

    "$0" umount "$DIR"
    sudo rm -rf "$DIR"
    sudo mv "$TEMP_DIR" "$DIR"
    sudo rm -rf "$OVERLAY_ROOT" "${OVERLAY_ROOT}.img"

    echo "Success: SquashFS image destroyed and data restored to \"$DIR\"."
    ;;

  *)
    usage
    ;;
esac