#!/bin/bash

# Arch Live Remaster
# Adapted from antiX live-remaster for Arch Linux
#
# Copyright (C) 2026 Adrian <adrian@mxlinux.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

# Default compression
COMP_TYPE="zstd"

# Version
VERSION="26.01"

# Default settings
BACKUP_ORIG=
SKIP_MD5=false
CREATED_MP=false
SAFETY_MARGIN_MIB=20
BOOT_MP=
OUTPUT_DIR=
LEAVE_NEW_FILE=false
EXCLUDE_FILE=

# Check for non-elevated options first
case "$1" in
    -h|--help)
        echo "Usage: $0 [-c|--comp compression_type] [-g|l|x|z] [-b|--no-backup] [-m|--no-checksum|--no-md5sum] [-v|--version] [-h|--help]"
        echo "Supported compression: gzip(-g), lz4(-l), xz(-x), zstd(-z)"
        exit 0
        ;;
    -v|--version)
        echo "Arch Live Remaster v$VERSION"
        exit 0
        ;;
esac

# Check for root privileges and self-elevate if needed
if [ "$(id -u)" -ne 0 ]; then
    exec sudo "$0" "$@"
fi

# Trap for cleanup
cleanup() {
    # Check for orphaned temporary backup (indicates interrupted replacement)
    if [ -n "$BOOT_MP" ] && [ -f "$BOOT_MP/arch/x86_64/airootfs.sfs.tmp.backup" ]; then
        echo ""
        echo "WARNING: Temporary backup file found at $BOOT_MP/arch/x86_64/airootfs.sfs.tmp.backup"
        echo "This may indicate an interrupted replacement. To recover:"
        echo "  mv $BOOT_MP/arch/x86_64/airootfs.sfs.tmp.backup $BOOT_MP/arch/x86_64/airootfs.sfs"
    fi

    if [ "$CREATED_MP" = true ]; then
        # Safety check: Prevent deletion of critical system paths
        case "$BOOT_MP" in
            "/"|"/home"|"/usr"|"/var"|"/etc"|"/boot"|"/root"|"/mnt"|"/media"|"/sys"|"/proc"|"/dev")
                echo "ERROR: Refusing to delete critical system path: $BOOT_MP"
                return 1
                ;;
        esac
        if umount "$BOOT_MP" 2>/dev/null; then
            rm -rf "$BOOT_MP"
        else
            echo "Failed to umount $BOOT_MP, skipping rm"
        fi
    fi
    [ -n "$OUTPUT_DIR" ] && [ "$LEAVE_NEW_FILE" != true ] && rm -f "$OUTPUT_DIR/airootfs.sfs.new"
    [ -n "$EXCLUDE_FILE" ] && rm -f "$EXCLUDE_FILE"
}
trap cleanup EXIT

# Function to get size in MiB
fmt_size() {
    local size=$1
    echo "$size MiB"
}

# Function to get free space in MiB
free_space() {
    df -BM "$1" | awk 'NR==2 {print $4}' | sed 's/M$//' | awk '{print int($1)}'
}

# Function to get total space in MiB
all_space() {
    df -BM "$1" | awk 'NR==2 {print $2}' | sed 's/M$//' | awk '{print int($1)}'
}

# Function to get du size in MiB
du_size() {
    du -sm "$1" 2>/dev/null | awk '{print int($1)}' || echo 0
}

# Function to remount a mount point read-write
remount_rw() {
    local mp=$1
    if mount -o remount,rw "$mp" 2>/dev/null; then
        echo "Remounted $mp RW"
        return 0
    else
        echo "RW remount failed for $mp"
        return 1
    fi
}

# Function to find and mount a live boot partition containing arch/x86_64
find_live_partition() {
    local label
    local dev
    local devices=""
    local tmp_mount="/tmp/archiso_boot"

    # Common labels seen on Arch-based live media
    for label in MXarch-Live ARCHISO_EFI LIVE-UEFI; do
        for dev in $(blkid -l -t LABEL="$label" -o device 2>/dev/null); do
            case " $devices " in
                *" $dev "*) ;;
                *) devices="$devices $dev" ;;
            esac
        done
    done

    # Official Arch ISO labels look like ARCH_YYYYMM
    for dev in $(blkid 2>/dev/null | awk -F: '/LABEL="ARCH_[0-9]{6}"/ {print $1}'); do
        case " $devices " in
            *" $dev "*) ;;
            *) devices="$devices $dev" ;;
        esac
    done

    mkdir -p "$tmp_mount"

    for dev in $devices; do
        if mount "$dev" "$tmp_mount" 2>/dev/null; then
            if [ -d "$tmp_mount/arch/x86_64" ]; then
                BOOT_MP="$tmp_mount"
                CREATED_MP=true
                echo "Mounted live boot partition ($dev) at $BOOT_MP"
                return 0
            fi
            umount "$tmp_mount" 2>/dev/null || true
        fi
    done

    rm -rf "$tmp_mount"
    return 1
}

# Function to check for required dependencies
check_dependencies() {
    local missing=""

    # Check for mksquashfs
    if ! command -v mksquashfs >/dev/null 2>&1; then
        missing="$missing mksquashfs"
    fi

    # Check for sha512sum (only if checksum is not skipped)
    if [ "$SKIP_MD5" = false ] && ! command -v sha512sum >/dev/null 2>&1; then
        missing="$missing sha512sum"
    fi

    if [ -n "$missing" ]; then
        echo "ERROR: Missing required tools:$missing"
        echo "Please install the required packages and try again."
        exit 1
    fi
}

# Main function
main() {
    echo "Arch Live Remaster v$VERSION"

    # Check dependencies
    check_dependencies

    # Excludes: runtime mounts and cache
    excludes="proc sys dev tmp mnt media lost+found var/tmp var/log var/cache var/lib/pacman/sync"

    # Set boot mount point
    BOOT_MP="/run/archiso/bootmnt"

    # Check for copytoram mode and handle boot partition mounting
    if [ -f "/run/archiso/copytoram/airootfs.sfs" ]; then
        echo "Detected copytoram mode. Attempting to mount boot partition for writing."
        if mountpoint -q "$BOOT_MP" 2>/dev/null && [ -d "$BOOT_MP/arch/x86_64" ]; then
            echo "Using existing boot mount point: $BOOT_MP"
        elif ! find_live_partition; then
            echo "Could not auto-detect a live boot partition."
            echo "Tried labels: ARCH_YYYYMM, ARCHISO_EFI, LIVE-UEFI, MXarch-Live"
            echo "Partition must contain arch/x86_64/"
        else
            echo "Auto-detected boot partition at $BOOT_MP"
        fi

        if [ "$BOOT_MP" = "/run/archiso/bootmnt" ]; then
            # Ask user to provide mount point
            read -rp "Please mount the live USB partition containing arch/x86_64 and provide the mount path: " USER_MP
            if [ -d "$USER_MP" ] && [ -d "$USER_MP/arch" ] && [ -d "$USER_MP/arch/x86_64" ]; then
                BOOT_MP="$USER_MP"
                echo "Using user-provided mount point: $BOOT_MP"
            else
                echo "ERROR: Invalid mount point provided. Directory must contain arch/x86_64/"
                exit 1
            fi
        fi
    fi

    # Check if boot mount point exists and is accessible
    if [ ! -d "$BOOT_MP" ]; then
        echo "ERROR: Boot mount point $BOOT_MP does not exist"
        echo "This script must be run from a live Arch ISO environment"
        exit 1
    fi

    # Check if boot mount point is actually a mountpoint (separate filesystem)
    if ! mountpoint -q "$BOOT_MP" 2>/dev/null; then
        echo "ERROR: $BOOT_MP is not a mounted filesystem"
        echo "This script must be run from a live Arch ISO environment where the boot device is mounted"
        exit 1
    fi

    # Verify boot mount point is on a different device than root
    root_device=$(df --output=source / | tail -n1)
    boot_device=$(df --output=source "$BOOT_MP" 2>/dev/null | tail -n1)

    if [ "$root_device" = "$boot_device" ] && [ -n "$boot_device" ]; then
        echo "WARNING: Boot mount point appears to be on the same device as root ($root_device)"
        echo "This is unusual for a live environment. Proceeding with caution..."
    else
        echo "Boot device detected: $boot_device"
    fi

    # Check if original airootfs.sfs exists
    SQFILE_FULL="$BOOT_MP/arch/x86_64/airootfs.sfs"
    if [ -f "$SQFILE_FULL" ]; then
        ORIG_EXISTS=true
        echo "Found existing airootfs.sfs at $SQFILE_FULL"
        # Ask about backup preference only if file exists and backup flag not set
        if [ -z "$BACKUP_ORIG" ]; then
            while true; do
                if ! read -rp "Backup original airootfs.sfs before replacing? (Y/n): " BACKUP_CHOICE; then
                    echo "EOF detected, defaulting to no backup."
                    BACKUP_ORIG=false
                    break
                fi
                case "$BACKUP_CHOICE" in
                    ""|y|Y)
                        BACKUP_ORIG=true
                        break
                        ;;
                    n|N)
                        BACKUP_ORIG=false
                        break
                        ;;
                    *)
                        echo "Please answer 'y' or 'n'"
                        ;;
                esac
            done
        fi
    else
        ORIG_EXISTS=false
        BACKUP_ORIG=false
        echo "No existing airootfs.sfs found (may have been deleted to save space)"
        echo "Will create new airootfs.sfs at $SQFILE_FULL"
    fi

    # Try to remount boot mount point as read-write
    if remount_rw "$BOOT_MP"; then
        RW_SUCCESS=true
    else
        RW_SUCCESS=false
    fi

    # Path to airootfs directory
    SQFILE_DIR="$BOOT_MP/arch/x86_64"
    OUTPUT_DIR="$SQFILE_DIR"

    # Ensure output directory exists
    mkdir -p "$SQFILE_DIR"

    # Check if device is writable
    if [ "$RW_SUCCESS" = false ]; then
        echo "ERROR: Boot device is read-only and cannot be remounted"
        echo "Please remount the device as read-write and run again"
        exit 1
    fi

    # Verify write access
    if ! touch "$BOOT_MP/.write_test" 2>/dev/null || ! rm "$BOOT_MP/.write_test" 2>/dev/null; then
        echo "ERROR: Boot device is not writable"
        echo "Check permissions and filesystem status"
        exit 1
    fi

    # Get sizes - calculate uncompressed size only for included directories
    # Excluding: proc sys dev tmp mnt media lost+found var/tmp var/log var/cache var/lib/pacman/sync
    # Also exclude mounted filesystems that shouldn't be part of rootfs
    fs_new_size=0
    for dir in /*; do
        dir_name="${dir#/}"
        # Skip excluded directories
        skip=false
        for excl in $excludes; do
            if [ "$dir_name" = "$excl" ]; then
                skip=true
                break
            fi
        done
        if [ "$skip" = false ] && [ -d "/$dir_name" ]; then
            # Skip /run: it is a tmpfs, so -one-file-system excludes it from the squashfs
            if [ "$dir_name" != "run" ]; then
                dir_size=$(du_size "/$dir_name")
                fs_new_size=$((fs_new_size + dir_size))
            fi
        fi
    done

    if [ "$ORIG_EXISTS" = true ]; then
        sq_old_size=$(du_size "$SQFILE_FULL")
        echo "Original airootfs size: $(fmt_size "$sq_old_size")"
    else
        sq_old_size=0
        echo "No original airootfs to compare"
    fi
    echo "Estimated new size (uncompressed): $(fmt_size $fs_new_size)"

    boot_total=$(all_space "$BOOT_MP")
    boot_avail=$(free_space "$BOOT_MP")

    echo "Boot device total: $(fmt_size "$boot_total")"
    echo "Boot device free: $(fmt_size "$boot_avail")"

    echo "Using compression: $COMP_TYPE"

    # Create exclude file
    EXCLUDE_FILE=$(mktemp /tmp/remaster.exclude.XXXXXX) || {
        echo "ERROR: Failed to create temporary exclude file"
        exit 1
    }
    {
        for excl in $excludes; do
            echo "$excl"
        done
        # Exclude runtime /run content from the live session
        echo "run/*"
    } > "$EXCLUDE_FILE"

    mkdir -p "$OUTPUT_DIR"
    # Temp file
    tmp_file="$OUTPUT_DIR/airootfs.sfs.tmp"
    new_file="$OUTPUT_DIR/airootfs.sfs.new"

    # Run mksquashfs
    echo "Creating new airootfs.sfs... This may take a while."
    if mksquashfs / "$tmp_file" \
        -comp "$COMP_TYPE" \
        -wildcards \
        -one-file-system \
        -ef "$EXCLUDE_FILE" \
        -p "run/archiso d 755 0 0" \
        -p "run/archiso/bootmnt d 755 0 0"; then
        echo "mksquashfs succeeded"
    else
        echo "mksquashfs failed"
        rm -f "$tmp_file"
        exit 1
    fi

    # Move to new
    echo "Moving to final file..."
    mv "$tmp_file" "$new_file"

    sq_new_size=$(du_size "$new_file")
    echo "New airootfs created: $new_file ($(fmt_size "$sq_new_size"))"

    # Enforce safety margin before replacing the original image
    boot_avail_after_build=$(free_space "$BOOT_MP")
    if [ "$boot_avail_after_build" -lt "$SAFETY_MARGIN_MIB" ]; then
        echo "WARNING: Insufficient free space for safe in-place replacement."
        echo "Free space: $(fmt_size "$boot_avail_after_build"), required safety margin: $(fmt_size "$SAFETY_MARGIN_MIB")"
        echo "Leaving new image at $new_file for manual replacement."
        echo "Manual steps:"
        echo "  1) Replace $SQFILE_FULL with $new_file"
        echo "  2) Regenerate checksum: (cd $SQFILE_DIR && sha512sum airootfs.sfs > airootfs.sha512)"
        LEAVE_NEW_FILE=true
        exit 0
    fi

    # Handle replacement based on backup preference
    echo "Preparing to replace airootfs.sfs..."
    if remount_rw "$BOOT_MP"; then
        # Create temporary backup for rollback protection
        tmp_backup="$SQFILE_FULL.tmp.backup"
        if [ "$ORIG_EXISTS" = true ]; then
            echo "Creating temporary backup for rollback protection..."
            if ! mv "$SQFILE_FULL" "$tmp_backup"; then
                echo "ERROR: Failed to create temporary backup"
                exit 1
            fi
        fi

        # Attempt replacement
        echo "Replacing with new airootfs.sfs..."
        if ! mv "$new_file" "$SQFILE_FULL"; then
            echo "ERROR: Failed to move new file into place"
            # Rollback: restore from temporary backup
            if [ -f "$tmp_backup" ]; then
                echo "Rolling back: restoring original airootfs.sfs..."
                mv "$tmp_backup" "$SQFILE_FULL"
            fi
            exit 1
        fi

        # Sync to ensure file is written to disk
        sync

        checksum_file="$SQFILE_DIR/airootfs.sha512"

        # Handle checksum
        if [ "$SKIP_MD5" = false ]; then
            echo "Updating SHA-512 checksum..."
            if ! (cd "$SQFILE_DIR" && sha512sum "airootfs.sfs" > "airootfs.sha512"); then
                echo "WARNING: SHA-512 calculation failed, but file replacement succeeded"
            fi
        else
            # Remove old checksum files if they exist to avoid confusion
            if [ -f "$checksum_file" ] || [ -f "$SQFILE_FULL.md5" ]; then
                echo "Removing old checksum file(s)..."
                rm -f "$checksum_file" "$SQFILE_FULL.md5"
            fi
        fi

        # Replacement succeeded - handle the backup based on user preference
        if [ -f "$tmp_backup" ]; then
            if [ "$BACKUP_ORIG" = true ]; then
                echo "Keeping backup as airootfs.sfs.old..."
                mv "$tmp_backup" "$SQFILE_FULL.old"
            else
                echo "Removing temporary backup..."
                rm -f "$tmp_backup"
            fi
        fi

        if [ "$BOOT_MP" != "/" ]; then
            # Verify $BOOT_MP is not on the root filesystem to avoid locking up the system
            # If both are on the same device, remounting RO would freeze the entire system
            root_device=$(df --output=source / | tail -n1)
            boot_device=$(df --output=source "$BOOT_MP" 2>/dev/null | tail -n1)
            # Safe to remount RO only if: different devices OR df failed to detect boot device
            if [ "$root_device" != "$boot_device" ] || [ -z "$boot_device" ]; then
                echo "Remounting $BOOT_MP RO..."
                mount -o remount,ro "$BOOT_MP"
            else
                echo "WARNING: Skipping RO remount (boot on same device as root: $root_device)"
            fi
        fi
        echo "Replacement complete. Reboot to test the new system."
    else
        echo "ERROR: Failed to remount $BOOT_MP RW for replacement"
        echo "Please manually replace: cp $new_file $SQFILE_FULL && (cd $SQFILE_DIR && sha512sum airootfs.sfs > airootfs.sha512)"
        exit 1
    fi
}

# Parse options
while [ $# -gt 0 ]; do
    case $1 in
        -h|--help)
            echo "Usage: $0 [-c|--comp compression_type] [-g|l|x|z] [-b|--no-backup] [-m|--no-checksum|--no-md5sum] [-v|--version] [-h|--help]"
            echo "Supported compression: gzip(-g), lz4(-l), xz(-x), zstd(-z)"
            exit 0
            ;;
        -v|--version)
            echo "Arch Live Remaster v$VERSION"
            exit 0
            ;;
        -c|--comp)
            if [ $# -lt 2 ] || [ -z "$2" ] || [[ "$2" == -* ]]; then
                echo "ERROR: --comp requires an argument. Supported: gzip, lz4, xz, zstd"
                exit 1
            fi
            COMP_TYPE="$2"
            shift 2
            # Validate compression type
            valid_comps="gzip lz4 xz zstd"
            if [[ ! " $valid_comps " =~ \ $COMP_TYPE\  ]]; then
                echo "ERROR: Unsupported compression type: $COMP_TYPE. Supported: $valid_comps"
                exit 1
            fi
            ;;
        --no-backup)
            BACKUP_ORIG=false
            shift
            ;;
        --no-checksum|--no-md5sum)
            SKIP_MD5=true
            shift
            ;;
        -*)
            OPTS=${1#-}
            COMP_SET_IN_THIS_ARG=false
            while [ -n "$OPTS" ]; do
                OPT_CHAR="${OPTS:0:1}"
                case "$OPT_CHAR" in
                    g|l|x|z)
                        if [ "$COMP_SET_IN_THIS_ARG" = true ]; then
                            echo "ERROR: Multiple compression types specified in same option group"
                            echo "Use only one of: -g, -l, -x, -z"
                            exit 1
                        fi
                        case "$OPT_CHAR" in
                            g) COMP_TYPE=gzip ;;
                            l) COMP_TYPE=lz4 ;;
                            x) COMP_TYPE=xz ;;
                            z) COMP_TYPE=zstd ;;
                        esac
                        COMP_SET_IN_THIS_ARG=true
                        ;;
                    b)
                        BACKUP_ORIG=false
                        ;;
                    m)
                        SKIP_MD5=true
                        ;;
                    *)
                        echo "Usage: $0 [-c|--comp compression_type] [-g|l|x|z] [-b|--no-backup] [-m|--no-checksum|--no-md5sum] [-v|--version] [-h|--help]"
                        echo "Supported compression: gzip(-g), lz4(-l), xz(-x), zstd(-z)"
                        exit 1
                        ;;
                esac
                OPTS=${OPTS#?}
            done
            shift
            ;;
        *)
            echo "Usage: $0 [-c|--comp compression_type] [-g|l|x|z] [-b|--no-backup] [-m|--no-checksum|--no-md5sum] [-v|--version] [-h|--help]"
            echo "Supported compression: gzip(-g), lz4(-l), xz(-x), zstd(-z)"
            exit 1
            ;;
    esac
done

main
