#!/bin/bash
# Author: Tomas M. <http://www.linux-live.org>
# Author: crims0n. <https://minios.dev>

VERSION="1.5"
SYSTEMNAME="MiniOS"
export TEXTDOMAIN="minios-tools"
TEMP=/tmp/minios.iso.$$
REGEX='^$'
MODULES=()
PWD="$(pwd)"
VERBOSITY_LEVEL=0
FORCE_BOOTLOADER=""
if [ -d "/run/initramfs/memory" ]; then
    SOURCE="/run/initramfs/memory"
elif [ -d "/memory" ]; then
    SOURCE="/memory"
fi

help() {
    echo "$(gettext "Usage"): $(basename $0) [OPTIONS]... [MODULE.SB]..."
    echo "$(gettext "Generate MiniOS ISO image, adding specified modules.")"
    echo ""
    echo "$(gettext "Options"):"
    echo "  -b, --bootloader TYPE $(gettext "Force specific BIOS bootloader type (isolinux or grub)")"
    echo "                        $(gettext "Note: EFI boot always uses GRUB regardless of this setting")"
    echo "  -e, --exclude REGEX   $(gettext "Exclude any existing path or file matching REGEX")"
    echo "  -n, --name NAME       $(gettext "Specify output ISO filename (default"): minios-YYYYMMDD_HHMM.iso)"
    echo "  -v, --verbose         $(gettext "Increase verbosity level")"
    echo "      --help            $(gettext "Display this help and exit")"
    echo "      --version         $(gettext "Display version information and exit")"
    echo ""
    echo "$(gettext "Bootloader Notes"):"
    echo "  - $(gettext "BIOS boot: Uses GRUB by default, can be forced to ISOLINUX/SYSLINUX with -b")"
    echo "  - $(gettext "EFI boot: Always uses GRUB (both 32-bit and 64-bit EFI support)")"
    echo "  - $(gettext "Hybrid ISOs support both BIOS and EFI boot modes simultaneously")"
    echo ""
    echo "$(gettext "Examples"):"
    echo "  # $(gettext "Create MiniOS ISO without firefox.sb module:")"
    echo "  $(basename $0) -e 'firefox' -f minios_without_firefox.iso"
    echo "  # $(gettext "Create MiniOS text-mode core only:")"
    echo "  $(basename $0) --exclude='firmware|xorg|desktop|apps|firefox' --file=minios_textmode.iso"
    echo "  # $(gettext "Force ISOLINUX bootloader for BIOS (EFI still uses GRUB):")"
    echo "  $(basename $0) --bootloader isolinux -n minios_isolinux.iso"
    echo "  # $(gettext "Use default GRUB bootloader for both BIOS and EFI:")"
    echo "  $(basename $0) -n minios_grub.iso"
    exit 0
}

brief_help() {
    echo "$(gettext "Usage"): $(basename $0) [OPTIONS]... [MODULE.SB]..."
    echo "$(gettext "Try") '$(basename $0) --help' $(gettext "for more information.")"
    exit 1
}

version() {
    echo "$(basename $0) $VERSION"
    exit 0
}

console_colors() {
    RED=$'\e[31m'
    GREEN=$'\e[32m'
    YELLOW=$'\e[33m'
    CYAN=$'\e[36m'

    BOLD=$'\e[1m'

    ENDCOLOR=$'\e[0m'
}

error() {
    local MESSAGE="${1-}"
    if [ "$OUTPUT_MODE" = "console" ]; then
        echo -e "${BOLD}${RED}E:${ENDCOLOR} $MESSAGE" >&2
    else
        echo "E: $MESSAGE" >&2
    fi
}

warning() {
    local MESSAGE="${1-}"
    if [ "$OUTPUT_MODE" = "console" ]; then
        echo -e "${BOLD}${YELLOW}W:${ENDCOLOR} $MESSAGE"
    else
        echo "W: $MESSAGE"
    fi
}

information() {
    local MESSAGE="${1-}"
    if [ "$OUTPUT_MODE" = "console" ]; then
        echo -e "${BOLD}${CYAN}I:${ENDCOLOR} $MESSAGE"
    else
        echo "I: $MESSAGE"
    fi
}

spinner() {
    local PID="${1}"
    local MSG="${2}"
    local XTRACE_WAS_SET=false
    case "$-" in
    *x*)
        XTRACE_WAS_SET=true
        set +x
        ;;
    esac

    local DELAY=0.1
    local SPINSTR='|/-\\'
    while [ -d "/proc/${PID}" ]; do
        for ((i = 0; i < ${#SPINSTR}; i++)); do
            printf "\r${BOLD}${CYAN}I:${ENDCOLOR} ${MSG} [${CYAN}${SPINSTR:$i:1}${ENDCOLOR}]"
            sleep "${DELAY}"
        done
    done
    printf "\r${BOLD}${CYAN}I:${ENDCOLOR} ${MSG} [${GREEN}done${ENDCOLOR}]$(tput el)\n"

    if ${XTRACE_WAS_SET}; then
        set -x
    fi
}

run_with_spinner() {
    local MSG="$1"
    shift
    if [ "${VERBOSITY_LEVEL}" -ge 1 ]; then
        "$@"
    else
        "$@" >/dev/null 2>&1 &
        local CMD_PID="$!"
        spinner "${CMD_PID}" "${MSG}"
        wait "${CMD_PID}"
    fi
}

console_colors

while [ "$#" -gt 0 ]; do
    case "$1" in
    -b | --bootloader)
        if [ "$2" = "isolinux" ] || [ "$2" = "grub" ]; then
            FORCE_BOOTLOADER="$2"
            shift 2
        else
            error "$(gettext "Invalid BIOS bootloader type: $2. Use 'isolinux' or 'grub'.")"
            information "$(gettext "Note: This option only affects BIOS boot. EFI boot always uses GRUB.")"
            brief_help
        fi
        ;;
    -e | --exclude)
        REGEX="$2"
        shift 2
        ;;
    -n | --name)
        TARGET="$(readlink -f "$2")"
        shift 2
        ;;
    -v | --verbose)
        VERBOSITY_LEVEL=$((VERBOSITY_LEVEL + 1))
        shift
        ;;
    --help)
        help
        ;;
    --version)
        version
        ;;
    *) # unknown option
        if [[ $1 == -* ]]; then
            brief_help
        elif [ -e "$1" ]; then
            MODULES+=("$(readlink -f "$1")")
            shift
        else
            error "$(gettext "Unknown argument: $1")"
            brief_help
        fi
        ;;
    esac
done

TARGET=${TARGET:-"${PWD}/minios-$(date +%Y%m%d_%H%M).iso"}

if [ -e "$SOURCE/data/minios/boot/isolinux.bin" ]; then
    MINIOS=$SOURCE/data/minios
elif [ -e "$SOURCE/data/minios/boot/grub/grub.cfg" ]; then
    MINIOS=$SOURCE/data/minios
elif [ -e "$SOURCE/iso/minios/boot/isolinux.bin" ]; then
    MINIOS=$SOURCE/iso/minios
elif [ -e "$SOURCE/iso/minios/boot/grub/grub.cfg" ]; then
    MINIOS=$SOURCE/iso/minios
elif [ -e "$SOURCE/toram/boot/isolinux.bin" ]; then
    MINIOS=$SOURCE/toram
elif [ -e "$SOURCE/toram/boot/grub/grub.cfg" ]; then
    MINIOS=$SOURCE/toram
elif [ -e "$SOURCE/data/from/0/minios/boot/isolinux.bin" ]; then
    MINIOS=$SOURCE/data/from/0/minios
elif [ -e "$SOURCE/data/from/0/minios/boot/grub/grub.cfg" ]; then
    MINIOS=$SOURCE/data/from/0/minios
fi

if [ "$MINIOS" = "" ]; then
    echo "Cannot find boot/isolinux.bin or boot/grub/grub.cfg in MiniOS data" >&2
    exit 2
fi

echo $SYSTEMNAME >/tmp/info

# Check required files
# Detect bootloader type
if [ "$FORCE_BOOTLOADER" != "" ]; then
    BOOTLOADER_TYPE="$FORCE_BOOTLOADER"
    information "$(gettext "Using forced BIOS bootloader type: $BOOTLOADER_TYPE")"
    information "$(gettext "Note: EFI boot will still use GRUB regardless of this setting")"
else
    # Default to GRUB, fallback to isolinux if GRUB files are missing
    if [ -e "$MINIOS/boot/grub/grub.cfg" ] && [ -e "$MINIOS/boot/grub/i386-pc/eltorito.img" ]; then
        BOOTLOADER_TYPE="grub"
        information "$(gettext "Using default GRUB bootloader for BIOS boot")"
    elif [ -e "$MINIOS/boot/isolinux.bin" ]; then
        BOOTLOADER_TYPE="isolinux"
        information "$(gettext "GRUB files not found, falling back to ISOLINUX for BIOS boot")"
    else
        BOOTLOADER_TYPE="grub"  # Default fallback
        information "$(gettext "Using GRUB bootloader (default)")"
    fi
fi

if [ "$BOOTLOADER_TYPE" = "isolinux" ]; then
    REQUIRED_FILES=(
        "$MINIOS/boot/isolinux.bin"
        "$MINIOS/boot/isohdpfx.bin"
        "$MINIOS/boot/syslinux.cfg"
    )
else
    REQUIRED_FILES=(
        "$MINIOS/boot/grub/i386-pc/eltorito.img"
        "$MINIOS/boot/grub/grub.cfg"
    )
fi

# Always check for EFI files
REQUIRED_FILES+=(
    "$MINIOS/boot/grub/efi64.img"
    "$MINIOS/boot/grub/efi32.img"
)

for FILE in "${REQUIRED_FILES[@]}"; do
    if [ ! -e "$FILE" ]; then
        if [ "$FORCE_BOOTLOADER" != "" ]; then
            error "$(gettext "Required file for forced BIOS bootloader '$FORCE_BOOTLOADER' not found: $FILE")"
            information "$(gettext "The selected BIOS bootloader files are missing. Try without --bootloader option for automatic detection.")"
            information "$(gettext "Remember: This option only affects BIOS boot, EFI boot always uses GRUB.")"
        else
            error "$(gettext "Required file not found: $FILE")"
            information "$(gettext "If you loaded the system into RAM, use 'toram=full' in the kernel parameters.")"
        fi
        exit 2
    fi
done

GRAFT=$(
    cd "$MINIOS"
    find . -type f | sed -r "s:^[.]/::" | egrep -v "^changes/" | egrep -v "$REGEX" | while read LINE; do
            echo "minios/$LINE=$MINIOS/$LINE"
        done
    if [ -d "$MINIOS/boot/EFI" ]; then
        cd "$MINIOS/boot/EFI"
        find . -type f | sed -r "s:^[.]/::" | egrep -v "$REGEX" | while read LINE; do
                echo "EFI/$LINE=$MINIOS/boot/EFI/$LINE"
            done
    fi
    echo ".disk/info=/tmp/info"
    # Add config file
    if [ -e "/etc/live/config.conf" ]; then
        echo "minios/config.conf=/etc/live/config.conf"
    fi
    # Add config.conf.d directory if it exists
    if [ -d "/etc/live/config.conf.d" ]; then
        find /etc/live/config.conf.d -type f | while read LINE; do
            BASENAME=$(basename "$LINE")
            echo "minios/config.conf.d/$BASENAME=$LINE"
        done
    fi
)

# add all modules
for MOD in "${MODULES[@]}"; do
    if [ ! -e "$MOD" ]; then
        error "$(gettext "File does not exist: $MOD")"
        exit 3
    fi
    BAS="$(basename "$MOD")"
    if [[ $BAS =~ ^[0-9]{2}- ]]; then
        GRAFT="$GRAFT minios/$BAS=$MOD"
    else
        GRAFT="$GRAFT minios/modules/$BAS=$MOD"
    fi
done

run_with_spinner "$(gettext "Creating ISO work directory")" mkdir -p "$TEMP/minios/boot" "$TEMP/minios/modules" "$TEMP/minios/changes" "$TEMP/minios/scripts"
if [ "$BOOTLOADER_TYPE" = "isolinux" ]; then
    run_with_spinner "$(gettext "Copying boot files (isolinux.bin & isohdpfx.bin)")" sh -c 'cp "$0/boot/isolinux.bin" "$1" && cp "$0/boot/isohdpfx.bin" "$1"' "$MINIOS" "$TEMP/minios/boot"
else
    run_with_spinner "$(gettext "Copying GRUB boot files")" mkdir -p "$TEMP/minios/boot/grub/i386-pc"
fi

cd "$TEMP" || exit 4

PERCHIMG=$(mktemp --suffix=.img)
run_with_spinner "$(gettext "Creating temporary persistence image")" dd if=/dev/zero of="$PERCHIMG" bs=1 count=0 seek=128k
run_with_spinner "$(gettext "Formatting persistence image")" mkfs.ext2 -b 1024 -L resizeme "$PERCHIMG"
# Build xorriso command based on bootloader type
XORRISO_CMD="xorriso --as mkisofs -iso-level 3 -volid \"${SYSTEMNAME^^}\" -A \"$SYSTEMNAME\" -joliet -joliet-long -rational-rock"

if [ "$BOOTLOADER_TYPE" = "isolinux" ]; then
    # ISOLINUX/SYSLINUX boot
    XORRISO_CMD="$XORRISO_CMD -eltorito-boot \"minios/boot/isolinux.bin\" -eltorito-catalog \"minios/boot/isolinux.boot\" -no-emul-boot -boot-load-size 4 -boot-info-table"
    MBR_FILE="$MINIOS/boot/isohdpfx.bin"
else
    # GRUB boot
    XORRISO_CMD="$XORRISO_CMD -eltorito-boot \"minios/boot/grub/i386-pc/eltorito.img\" -no-emul-boot -boot-load-size 4 -boot-info-table"
    MBR_FILE="$MINIOS/boot/grub/i386-pc/boot_hybrid.img"
fi

# EFI boot (common for both)
XORRISO_CMD="$XORRISO_CMD -eltorito-alt-boot -e \"minios/boot/grub/efi64.img\" -no-emul-boot"
XORRISO_CMD="$XORRISO_CMD -eltorito-alt-boot -e \"minios/boot/grub/efi32.img\" -no-emul-boot"

# Hybrid and partitioning
if [ -e "$MBR_FILE" ]; then
    XORRISO_CMD="$XORRISO_CMD --isohybrid-mbr \"$MBR_FILE\""
fi
XORRISO_CMD="$XORRISO_CMD -append_partition 2 0x83 \"$PERCHIMG\" -partition_cyl_align on -partition_offset 16 -part_like_isohybrid"
XORRISO_CMD="$XORRISO_CMD -graft-points $GRAFT -output \"$TARGET\""

# Execute with appropriate verbosity
if [[ "${VERBOSITY_LEVEL}" -ge 1 ]]; then
    bash -c "$XORRISO_CMD" && information "The image ${CYAN}$TARGET${ENDCOLOR} has been created." || exit 1
else
    bash -c "$XORRISO_CMD" >/dev/null 2>&1 &
    iso_pid=$!
    spinner "${iso_pid}" "$(gettext "Generating ISO image")"
    wait "${iso_pid}"
    if [ $? -eq 0 ]; then
        information "The image ${CYAN}$TARGET${ENDCOLOR} has been created."
    else
        exit 1
    fi
fi

run_with_spinner "$(gettext "Cleaning up temporary files")" rm -rf "$TEMP" "$PERCHIMG"

# Generate SHA256 checksum
if [[ -f "$TARGET" ]]; then
    sha256sum "$TARGET" > "$TARGET.sha256"
    information "SHA256 checksum saved to ${CYAN}$TARGET.sha256${ENDCOLOR}."
fi

information "ISO successfully created: ${CYAN}$TARGET${ENDCOLOR}"
exit 0
