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