#!/bin/bash

# Update CoW space in GRUB config for Arch live USB
# Usage: ./update-cow-space
#
# 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/>.

# Version
VERSION="26.01"

set -euo pipefail

ME=${0##*/}

fatal() { echo "$ME: error: $*" >&2; exit 1; }

usage() {
    cat <<USAGE
Usage: $ME <cow_size> [options]

Update the CoW space parameter in GRUB config on the running Arch live USB.

Arguments:
  <cow_size>     CoW space size (e.g., 1G, 512M)

This script detects if running from live USB and updates the appropriate GRUB config.
For copytoram mode, it will attempt to mount the boot partition automatically.

Options:
  -h, --help     Show this help message
  -v, --version  Show version information
USAGE
}

# Parse options and arguments
COW_SIZE=""
ORIG_ARGS=("$@")
while [ $# -gt 0 ]; do
    case $1 in
        -h|--help) usage; exit 0 ;;
        -v|--version) echo "Update CoW Space v$VERSION"; exit 0 ;;
        -*)
            usage
            exit 1
            ;;
        *)
            if [ -z "$COW_SIZE" ]; then
                COW_SIZE=$1
            else
                usage
                exit 1
            fi
            ;;
    esac
    shift
done

if [ -z "$COW_SIZE" ]; then
    usage
    exit 1
fi

# Allow help without root privileges (already handled above)

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

# Check if running in live environment
[[ -d /run/archiso ]] || fatal "Not running in Arch live environment (/run/archiso not found)"

MOUNT_POINT=""
CREATED_MP=false
TMP_MOUNT=""

cleanup() {
    if [ "$CREATED_MP" = true ] && [ -n "$MOUNT_POINT" ]; then
        if mountpoint -q "$MOUNT_POINT" 2>/dev/null; then
            if umount "$MOUNT_POINT" 2>/dev/null; then
                [ -n "$TMP_MOUNT" ] && [ -d "$TMP_MOUNT" ] && rmdir "$TMP_MOUNT" 2>/dev/null || true
                echo "Unmounted and cleaned up temporary mount"
            else
                echo "Failed to unmount $MOUNT_POINT"
            fi
        elif [ -n "$TMP_MOUNT" ] && [ -d "$TMP_MOUNT" ]; then
            rmdir "$TMP_MOUNT" 2>/dev/null || true
        fi
    fi
}
trap cleanup EXIT

# Find the ESP partition (EFI System Partition) by label
ESP_LABEL="LIVE-UEFI"
ESP_PART=$(blkid -l -t LABEL="$ESP_LABEL" -o device 2>/dev/null)
if [ -z "$ESP_PART" ]; then
    # Fallback to ARCHISO_EFI
    ESP_PART=$(blkid -l -t LABEL=ARCHISO_EFI -o device 2>/dev/null)
fi
if [ -z "$ESP_PART" ]; then
    ESP_PART=$(blkid -l -t LABEL=MXarch-Live -o device 2>/dev/null)
fi
if [ -z "$ESP_PART" ]; then
    fatal "Could not find ESP partition with labels LIVE-UEFI, ARCHISO_EFI, or MXarch-Live"
fi

# Check if bootmnt is already mounted to the ESP
BOOTMNT_DEVICE=$(findmnt -n -o SOURCE /run/archiso/bootmnt 2>/dev/null)
if [ "$BOOTMNT_DEVICE" = "$ESP_PART" ] && [ -d /run/archiso/bootmnt ]; then
    # bootmnt is correctly mounted to ESP
    MOUNT_POINT="/run/archiso/bootmnt"
    echo "Using existing bootmnt at $MOUNT_POINT (mounted to ESP $ESP_PART)"
else
    # Need to mount ESP
    TMP_MOUNT="/tmp/archiso_esp"
    mkdir -p "$TMP_MOUNT"
    if mount "$ESP_PART" "$TMP_MOUNT" 2>/dev/null; then
        MOUNT_POINT="$TMP_MOUNT"
        CREATED_MP=true
        echo "Mounted ESP ($ESP_PART) at $MOUNT_POINT"
    else
        echo "Failed to mount ESP $ESP_PART"
        # Ask user to provide mount point
        read -rp "Please mount the ESP ($ESP_PART) and provide the mount path: " USER_MP
        if [ -d "$USER_MP" ] && [ -d "$USER_MP/boot" ] && [ -f "$USER_MP/boot/grub/grub.cfg" ]; then
            MOUNT_POINT="$USER_MP"
            echo "Using user-provided mount point: $MOUNT_POINT"
        else
            fatal "Invalid mount point provided. Directory must contain boot/grub/grub.cfg"
        fi
    fi
fi

# Check if mount point exists and is accessible
[[ -d $MOUNT_POINT ]] || fatal "Mount point '$MOUNT_POINT' does not exist"
if ! mountpoint -q "$MOUNT_POINT" 2>/dev/null; then
    fatal "$MOUNT_POINT is not a mounted filesystem"
fi

# Check for GRUB config
GRUB_CFG="$MOUNT_POINT/boot/grub/grub.cfg"
[[ -f $GRUB_CFG ]] || fatal "GRUB config not found at $GRUB_CFG"

# Validate COW_SIZE
[[ $COW_SIZE =~ ^[0-9]+[MG]$ ]] || fatal "Invalid size format: $COW_SIZE. Use e.g., 1G or 512M"

# Check if read-write
if ! touch "$MOUNT_POINT/.test_write" 2>/dev/null; then
    echo "Partition is read-only, attempting to remount as RW..."
    # Find the device
    DEVICE=$(findmnt -n -o SOURCE "$MOUNT_POINT" 2>/dev/null)
    [[ -n $DEVICE ]] || fatal "Could not determine device for $MOUNT_POINT"
    mount -o remount,rw "$DEVICE" "$MOUNT_POINT" || fatal "Failed to remount $DEVICE as RW"
    echo "Remounted as RW"
else
    rm -f "$MOUNT_POINT/.test_write"
fi

# Backup GRUB config (unique per run)
BACKUP_GRUB_CFG="$GRUB_CFG.bak.$(date +%Y%m%d-%H%M%S).$$"
cp "$GRUB_CFG" "$BACKUP_GRUB_CFG"
echo "Backed up GRUB config to $BACKUP_GRUB_CFG"

# Update GRUB config only on boot linux lines.
# For each linux line: replace existing cow_spacesize, otherwise append it.
if ! sed -i "/^[[:space:]]*linux / { /cow_spacesize=/ s/cow_spacesize=[^ ]*/cow_spacesize=$COW_SIZE/; /cow_spacesize=/! s|$| cow_spacesize=$COW_SIZE|; }" "$GRUB_CFG"; then
    fatal "Failed to update cow_spacesize in GRUB config"
fi

# Verify the update
if ! grep -Eq "^[[:space:]]*linux .*[[:space:]]cow_spacesize=$COW_SIZE([[:space:]]|$)" "$GRUB_CFG"; then
    echo "Warning: cow_spacesize may not have been added correctly."
    echo "Current linux line(s):"
    grep -E "^[[:space:]]*linux " "$GRUB_CFG" || echo "No linux line found"
    fatal "Verification failed: cow_spacesize not found in config"
fi

echo "Updated GRUB config with cow_spacesize=$COW_SIZE"
echo "Reboot the USB to apply changes."
