#!/bin/bash
#===============================================================================
# PATCH OMARCHY ISO FOR DUAL-BOOT
#
# This version hooks into the original Omarchy installer (archinstall)
# We just handle partitioning, then let archinstall do everything else
#
# Usage: sudo ./patch-omarchy-dualboot.sh omarchy-3.x.x.iso
#===============================================================================

set -e

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

log() { echo -e "${GREEN}[✓]${NC} $1"; }
warn() { echo -e "${YELLOW}[!]${NC} $1"; }
error() { echo -e "${RED}[✗]${NC} $1"; exit 1; }

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"

[[ $EUID -ne 0 ]] && error "Run as root: sudo $0"

# Auto-detect ISO if not provided
if [[ -z "$1" ]]; then
    SOURCE_ISO=$(find "$SCRIPT_DIR" -maxdepth 1 -name "omarchy-*.iso" ! -name "omarchy-dualboot-*" ! -name "omarchy-installer-*" 2>/dev/null | head -1)
    if [[ -z "$SOURCE_ISO" ]]; then
        error "No Omarchy ISO found in $SCRIPT_DIR

Please either:
  1. Place omarchy-*.iso in the same directory as this script
  2. Or specify the path: sudo $0 /path/to/omarchy.iso"
    fi
    log "Auto-detected ISO: $(basename "$SOURCE_ISO")"
else
    SOURCE_ISO="$1"
fi

[[ ! -f "$SOURCE_ISO" ]] && error "ISO not found: $SOURCE_ISO"
WORK_DIR="$SCRIPT_DIR/.omarchy-patch-$$"
OUTPUT_ISO="$SCRIPT_DIR/omarchy-dualboot-$(date +%Y.%m.%d).iso"

# Check and install dependencies
MISSING_PKGS=""
command -v xorriso &>/dev/null || MISSING_PKGS="$MISSING_PKGS xorriso"
command -v unsquashfs &>/dev/null || MISSING_PKGS="$MISSING_PKGS squashfs-tools"
command -v isoinfo &>/dev/null || MISSING_PKGS="$MISSING_PKGS cdrtools"

if [[ -n "$MISSING_PKGS" ]]; then
    log "Installing missing dependencies:$MISSING_PKGS"
    pacman -S --noconfirm --needed $MISSING_PKGS || error "Failed to install dependencies"
    log "Dependencies installed"
fi

log "Patching: $SOURCE_ISO"

# Cleanup on exit
cleanup() {
    umount "$WORK_DIR/iso" 2>/dev/null || true
    rm -rf "$WORK_DIR"
}
trap cleanup EXIT

mkdir -p "$WORK_DIR"/{iso,extracted,newiso}

#===============================================================================
# EXTRACT ISO
#===============================================================================

log "Extracting ISO..."
mount -o loop,ro "$SOURCE_ISO" "$WORK_DIR/iso"
cp -a "$WORK_DIR/iso/"* "$WORK_DIR/newiso/"
umount "$WORK_DIR/iso"

# Find squashfs
SQUASHFS="$WORK_DIR/newiso/arch/x86_64/airootfs.sfs"
[[ ! -f "$SQUASHFS" ]] && error "Squashfs not found at expected path"

log "Extracting squashfs (this takes a few minutes)..."
unsquashfs -d "$WORK_DIR/extracted" "$SQUASHFS"

# Verify squashfs structure
if [[ ! -d "$WORK_DIR/extracted/root" ]]; then
    error "Unexpected squashfs structure: /root directory not found!"
fi

log "Original /root contents:"
ls -la "$WORK_DIR/extracted/root/" | grep -E '\.(zlogin|zprofile|automated|bash)' || echo "  (no matching files)"

#===============================================================================
# CREATE DUAL-BOOT WRAPPER SCRIPT
#===============================================================================

log "Creating dual-boot wrapper..."

# This script runs BEFORE the original installer
# It handles partitioning, then lets archinstall do the rest
cat > "$WORK_DIR/extracted/root/dualboot-setup.sh" << 'DUALBOOT_SCRIPT'
#!/bin/bash
#===============================================================================
# DUAL-BOOT PARTITION SETUP
# Creates partitions in free space, then hands off to archinstall
#===============================================================================

set -e

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'

log() { echo -e "${GREEN}[✓]${NC} $1"; }
warn() { echo -e "${YELLOW}[!]${NC} $1"; }
error() { echo -e "${RED}[✗]${NC} $1"; exit 1; }
info() { echo -e "${CYAN}[i]${NC} $1"; }

header() {
    echo ""
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}  $1${NC}"
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo ""
}

cleanup_and_exit() {
    warn "Setup cancelled or failed. Cleaning up..."
    umount -R /mnt 2>/dev/null || true
    cryptsetup close cryptroot 2>/dev/null || true
    exit 1
}

#===============================================================================
# CLEANUP FUNCTION FOR FAILED INSTALLATIONS
#===============================================================================

cleanup_failed_install() {
    header "CLEANUP FAILED INSTALLATION"

    echo "This will remove Linux partitions from a failed installation attempt."
    echo "Windows partitions will NOT be touched."
    echo ""

    echo "Available drives:"
    lsblk -d -o NAME,SIZE,MODEL | grep -E "^(nvme|sd|vd)"
    echo ""

    read -p "Enter drive to clean (e.g., nvme0n1): " drive_input
    DRIVE="/dev/$drive_input"

    [[ ! -b "$DRIVE" ]] && error "Drive not found: $DRIVE"
    [[ "$DRIVE" == *"nvme"* ]] && PART_PREFIX="${DRIVE}p" || PART_PREFIX="$DRIVE"

    echo ""
    echo "Current partitions:"
    lsblk -o NAME,SIZE,FSTYPE,LABEL "$DRIVE"
    echo ""

    cryptsetup close cryptroot 2>/dev/null || true
    umount -R /mnt 2>/dev/null || true

    LINUX_EFI=""
    LINUX_ROOT=""

    for part in ${PART_PREFIX}*; do
        [[ "$part" == "$DRIVE" || "$part" == "${PART_PREFIX}" ]] && continue
        [[ ! -b "$part" ]] && continue

        LABEL=$(lsblk -n -o LABEL "$part" 2>/dev/null | tr -d ' ')
        FSTYPE=$(lsblk -n -o FSTYPE "$part" 2>/dev/null | tr -d ' ')

        if [[ "$LABEL" == "LINUXEFI" ]]; then
            LINUX_EFI="$part"
            echo -e "${YELLOW}Found Linux EFI partition: $part${NC}"
        elif [[ "$FSTYPE" == "crypto_LUKS" ]]; then
            LINUX_ROOT="$part"
            echo -e "${YELLOW}Found LUKS partition: $part${NC}"
        fi
    done

    if [[ -z "$LINUX_EFI" && -z "$LINUX_ROOT" ]]; then
        echo ""
        info "No Linux partitions found from failed installation."
        read -p "Press Enter to continue..." _
        return 1
    fi

    echo ""
    echo -e "${RED}WARNING: The following partitions will be DELETED:${NC}"
    [[ -n "$LINUX_EFI" ]] && echo "  - $LINUX_EFI (Linux EFI)"
    [[ -n "$LINUX_ROOT" ]] && echo "  - $LINUX_ROOT (Linux Root/LUKS)"
    echo ""

    read -p "Type 'DELETE' to confirm removal: " confirm
    if [[ "$confirm" != "DELETE" ]]; then
        echo "Cancelled."
        return 1
    fi

    if [[ -n "$LINUX_ROOT" ]]; then
        PART_NUM=$(echo "$LINUX_ROOT" | grep -oE '[0-9]+$')
        log "Deleting partition $PART_NUM..."
        sgdisk -d "$PART_NUM" "$DRIVE"
    fi

    if [[ -n "$LINUX_EFI" ]]; then
        PART_NUM=$(echo "$LINUX_EFI" | grep -oE '[0-9]+$')
        log "Deleting partition $PART_NUM..."
        sgdisk -d "$PART_NUM" "$DRIVE"
    fi

    partprobe "$DRIVE"

    for entry in $(efibootmgr 2>/dev/null | grep -i "Omarchy" | sed -n 's/^Boot\([0-9A-Fa-f]*\).*/\1/p'); do
        efibootmgr -b "$entry" -B 2>/dev/null || true
    done

    log "Cleanup complete!"
    echo ""
    lsblk -o NAME,SIZE,FSTYPE,LABEL "$DRIVE"
    echo ""
    read -p "Press Enter to continue..." _
    return 0
}

#===============================================================================
# MAIN MENU
#===============================================================================

header "OMARCHY DUAL-BOOT INSTALLER"

echo "This installer will set up Omarchy alongside Windows."
echo "Windows partitions will NOT be touched."
echo ""
echo "Options:"
echo "  1) Install Omarchy (dual-boot with Windows)"
echo "  2) Clean up failed installation"
echo "  3) Standard install (wipe entire disk)"
echo "  4) Exit to shell"
echo ""
read -p "Select option [1]: " MENU_CHOICE
MENU_CHOICE=${MENU_CHOICE:-1}

case $MENU_CHOICE in
    2)
        cleanup_failed_install || true
        exec /root/dualboot-setup.sh
        ;;
    3)
        # Run original installer
        info "Starting standard Omarchy installer..."
        exec /root/.automated_script.sh.orig
        ;;
    4)
        echo "Type '/root/dualboot-setup.sh' to restart installer."
        exit 0
        ;;
    1)
        # Continue with dual-boot setup
        ;;
    *)
        exec /root/dualboot-setup.sh
        ;;
esac

#===============================================================================
# PRE-FLIGHT CHECKS
#===============================================================================

[[ ! -d /sys/firmware/efi ]] && error "UEFI mode required for dual-boot"

header "SELECT DRIVE"

echo "Available drives:"
echo ""
lsblk -d -o NAME,SIZE,MODEL | grep -E "^(nvme|sd|vd)"
echo ""

mapfile -t DRIVES < <(lsblk -d -n -o NAME,TRAN | awk '$2!="usb" {print $1}' | grep -E "^(nvme|sd|vd)")

if [[ ${#DRIVES[@]} -eq 1 ]]; then
    DRIVE="/dev/${DRIVES[0]}"
    info "Found: $DRIVE"
    read -p "Use this drive? (Y/n): " confirm
    [[ "$confirm" =~ ^[Nn]$ ]] && error "Cancelled"
else
    read -p "Enter drive (e.g., nvme0n1): " drive_input
    DRIVE="/dev/$drive_input"
fi

[[ ! -b "$DRIVE" ]] && error "Drive not found: $DRIVE"
[[ "$DRIVE" == *"nvme"* ]] && PART_PREFIX="${DRIVE}p" || PART_PREFIX="$DRIVE"

echo ""
echo "Current partitions on $DRIVE:"
lsblk -o NAME,SIZE,FSTYPE,LABEL "$DRIVE"
echo ""

# Check for Windows
if lsblk -o FSTYPE "$DRIVE" | grep -q ntfs; then
    log "Windows detected - it will be preserved"
else
    warn "Windows not detected on this drive"
    read -p "Continue anyway? (y/N): " cont
    [[ ! "$cont" =~ ^[Yy]$ ]] && error "Cancelled"
fi

# Check free space
FREE_SECTORS=$(sgdisk -p "$DRIVE" 2>/dev/null | grep "Total free space" | awk '{print $5}')
if [[ "$FREE_SECTORS" =~ ^[0-9]+$ ]]; then
    FREE_GB=$((FREE_SECTORS * 512 / 1024 / 1024 / 1024))
    [[ $FREE_GB -lt 20 ]] && error "Need 20GB+ free space, found ${FREE_GB}GB"
    log "Free space: ${FREE_GB}GB"
else
    warn "Could not determine free space"
fi

#===============================================================================
# ENCRYPTION PASSWORD
#===============================================================================

header "DISK ENCRYPTION"

echo "Your Linux partition will be encrypted with LUKS2."
echo "You'll enter your user password in the next step (configurator)."
echo ""

while true; do
    read -s -p "Enter LUKS encryption password: " LUKS_PASS; echo ""
    read -s -p "Confirm password: " LUKS_PASS2; echo ""
    [[ "$LUKS_PASS" == "$LUKS_PASS2" ]] && break
    warn "Passwords don't match, try again"
done

#===============================================================================
# CONFIRMATION
#===============================================================================

header "CONFIRM PARTITIONING"

echo "The following will be created in FREE SPACE on $DRIVE:"
echo ""
echo "  • 1GB EFI partition (for Linux bootloader)"
echo "  • Rest: LUKS2 encrypted partition (for Linux root)"
echo ""
echo -e "${GREEN}Windows partitions will NOT be touched.${NC}"
echo ""
read -p "Type 'yes' to continue: " confirm
[[ "$confirm" != "yes" ]] && error "Cancelled"

#===============================================================================
# CREATE PARTITIONS
#===============================================================================

header "CREATING PARTITIONS"

trap cleanup_and_exit ERR

DRIVE_NAME=$(basename "$DRIVE")
LAST_PART=$(lsblk -n -o NAME "$DRIVE" | grep -v "^${DRIVE_NAME}$" | grep -oE '[0-9]+$' | sort -n | tail -1)
LAST_PART=${LAST_PART:-0}

# Create EFI partition
log "Creating Linux EFI partition (1GB)..."
NEXT=$((LAST_PART + 1))
sgdisk -n "${NEXT}:0:+1G" -t "${NEXT}:ef00" "$DRIVE"
EFI_PART="${PART_PREFIX}${NEXT}"
partprobe "$DRIVE"; udevadm settle; sleep 2
mkfs.fat -F32 -n "LINUXEFI" "$EFI_PART"
log "Created: $EFI_PART"

# Create root partition
log "Creating Linux root partition..."
NEXT=$((NEXT + 1))
sgdisk -n "${NEXT}:0:0" -t "${NEXT}:8300" "$DRIVE"
ROOT_PART="${PART_PREFIX}${NEXT}"
partprobe "$DRIVE"; udevadm settle; sleep 2
log "Created: $ROOT_PART"

#===============================================================================
# SETUP ENCRYPTION & FILESYSTEM
#===============================================================================

header "SETTING UP ENCRYPTION"

log "Formatting LUKS2..."
echo -n "$LUKS_PASS" | cryptsetup luksFormat --type luks2 --key-file=- "$ROOT_PART"

log "Opening encrypted volume..."
echo -n "$LUKS_PASS" | cryptsetup open --key-file=- "$ROOT_PART" cryptroot
LUKS_UUID=$(blkid -s UUID -o value "$ROOT_PART")

log "Creating btrfs filesystem..."
mkfs.btrfs -f /dev/mapper/cryptroot

log "Creating subvolumes..."
mount /dev/mapper/cryptroot /mnt
btrfs subvolume create /mnt/@
btrfs subvolume create /mnt/@home
btrfs subvolume create /mnt/@log
btrfs subvolume create /mnt/@pkg
# NOTE: Do NOT create @snapshots - snapper will create .snapshots as nested subvolume
umount /mnt

log "Mounting filesystems..."
mount -o compress=zstd,subvol=@ /dev/mapper/cryptroot /mnt
mkdir -p /mnt/{home,var/log,var/cache/pacman/pkg,boot}
mount -o compress=zstd,subvol=@home /dev/mapper/cryptroot /mnt/home
mount -o compress=zstd,subvol=@log /dev/mapper/cryptroot /mnt/var/log
mount -o compress=zstd,subvol=@pkg /dev/mapper/cryptroot /mnt/var/cache/pacman/pkg
# NOTE: Do NOT mount .snapshots - snapper creates it as nested subvolume in @
mount "$EFI_PART" /mnt/boot

log "Partitions ready!"

# Save info for later
echo "$LUKS_UUID" > /tmp/luks_uuid
echo "$EFI_PART" > /tmp/efi_part
echo "$ROOT_PART" > /tmp/root_part
echo "$DRIVE" > /tmp/install_drive

trap - ERR

#===============================================================================
# HAND OFF TO ARCHINSTALL
#===============================================================================

header "STARTING OMARCHY INSTALLER"

echo "Partitions are ready. Now the Omarchy installer will collect your"
echo "user information and install the system."
echo ""
info "The disk is already partitioned - just enter your user details."
echo ""
read -p "Press Enter to continue to Omarchy setup..." _

# Run the configurator to get user info
cd /root
source /root/omarchy/install/helpers/all.sh 2>/dev/null || true

# Set colors
if [[ $(tty) == "/dev/tty"* ]]; then
    echo -en "\e]P01a1b26"
    echo -en "\e]P7a9b1d6"
    echo -en "\e]PFc0caf5"
    echo -en "\033[0m"
    clear
fi

# Run configurator
./configurator

# Get username from generated config
OMARCHY_USER=$(jq -r '.users[0].username' user_credentials.json 2>/dev/null)

if [[ -z "$OMARCHY_USER" || "$OMARCHY_USER" == "null" ]]; then
    error "Failed to get username from configurator"
fi

log "User: $OMARCHY_USER"

#===============================================================================
# MODIFY CONFIG FOR PRE-MOUNTED PARTITIONS
#===============================================================================

header "CONFIGURING INSTALLATION"

log "Updating configuration for dual-boot..."

# Only modify disk_config and remove bootloader, preserve everything else from generated config
# This keeps the correct kernel, locale, packages, profile, etc.
# We MUST remove bootloader because archinstall crashes when trying to
# detect root device with pre_mounted_config - it fails BEFORE installing packages
# We install limine manually after archinstall completes
jq '.disk_config = {"config_type": "pre_mounted_config", "mountpoint": "/mnt"} | del(.bootloader)' \
    user_configuration.json > user_configuration.json.tmp
mv user_configuration.json.tmp user_configuration.json

# Extract kernel name for later use in bootloader config
KERNEL_NAME=$(jq -r '.kernels[0] // "linux"' user_configuration.json)
echo "$KERNEL_NAME" > /tmp/kernel_name
log "Using kernel: $KERNEL_NAME"

log "Configuration updated for pre-mounted partitions"

#===============================================================================
# RUN ARCHINSTALL
#===============================================================================

header "INSTALLING SYSTEM"

log "Running archinstall..."
info "This will take several minutes..."

# Initialize keyring
pacman-key --init
pacman-key --populate archlinux
pacman-key --populate omarchy 2>/dev/null || true

# For offline install, ensure we use the offline-only pacman config
# and skip database sync since the database is already on the ISO
cat > /tmp/pacman-offline.conf << 'PACCONF'
[options]
HoldPkg     = pacman glibc
Architecture = auto
SigLevel    = Required DatabaseOptional
LocalFileSigLevel = Optional

[offline]
SigLevel = Optional TrustAll
Server = file:///var/cache/omarchy/mirror/offline/
PACCONF

# Overwrite system pacman.conf with offline-only version
cp /tmp/pacman-offline.conf /etc/pacman.conf

# Debug: verify config
log "Pacman config set to offline-only:"
grep -E "^\[|^Server" /etc/pacman.conf

# Skip pacman -Sy for offline install - database already exists on ISO
# Check for either offline.db or offline.db.tar.gz
if [[ -f /var/cache/omarchy/mirror/offline/offline.db ]] || [[ -f /var/cache/omarchy/mirror/offline/offline.db.tar.gz ]]; then
    log "Offline package database found - skipping sync"
else
    warn "Offline database not found - attempting sync..."
    pacman --config /tmp/pacman-offline.conf -Sy --noconfirm
fi

# Run archinstall with our config (same flags as original Omarchy installer)
# bootloader=none prevents archinstall from crashing on pre_mounted_config
archinstall \
    --config user_configuration.json \
    --creds user_credentials.json \
    --silent \
    --skip-ntp \
    --skip-wkd \
    --skip-wifi-check

# Verify archinstall created the user
if [[ ! -d /mnt/home/$OMARCHY_USER ]]; then
    error "archinstall failed - user $OMARCHY_USER not created"
fi

log "Base system installed!"

# Install limine bootloader package (archinstall didn't because we removed bootloader config)
log "Installing limine bootloader..."
# Copy offline pacman config to chroot and mount offline mirror first
cp /tmp/pacman-offline.conf /mnt/etc/pacman.conf
mkdir -p /mnt/var/cache/omarchy/mirror/offline
mount --bind /var/cache/omarchy/mirror/offline /mnt/var/cache/omarchy/mirror/offline
# Install limine using offline repo (no -Sy needed, packages are local)
arch-chroot /mnt pacman --noconfirm --needed -S limine

#===============================================================================
# POST-INSTALL CONFIGURATION
#===============================================================================

header "CONFIGURING BOOTLOADER"

LUKS_UUID=$(cat /tmp/luks_uuid)
KERNEL_NAME=$(cat /tmp/kernel_name)

# Configure mkinitcpio for encryption using Omarchy's drop-in approach
log "Configuring initramfs for encryption..."

# Install required packages
log "Installing plymouth and limine tools..."
arch-chroot /mnt pacman --noconfirm --needed -S plymouth limine-snapper-sync limine-mkinitcpio-hook || warn "Some packages may already be installed"

# Create Omarchy-style mkinitcpio drop-in config with encryption support
# This uses the same hooks as Omarchy but includes encrypt for LUKS
mkdir -p /mnt/etc/mkinitcpio.conf.d
cat > /mnt/etc/mkinitcpio.conf.d/omarchy_hooks.conf << 'MKINIT'
HOOKS=(base udev plymouth keyboard autodetect microcode modconf kms keymap consolefont block encrypt filesystems fsck btrfs-overlayfs)
MKINIT

cat > /mnt/etc/mkinitcpio.conf.d/thunderbolt_module.conf << 'MKINIT'
MODULES+=(thunderbolt)
MKINIT

# Rebuild initramfs (warnings about missing firmware are OK)
log "Building initramfs (warnings about missing firmware are normal)..."
if ! arch-chroot /mnt mkinitcpio -P; then
    warn "mkinitcpio reported warnings - checking if initramfs was created..."
fi

# Verify the initramfs was actually created for our kernel
if [[ -f /mnt/boot/initramfs-${KERNEL_NAME}.img ]]; then
    log "initramfs-${KERNEL_NAME}.img created successfully"
elif [[ -f /mnt/boot/initramfs-linux.img ]]; then
    log "initramfs-linux.img created successfully"
else
    error "initramfs was not created! Check mkinitcpio output above."
fi

# Verify kernel was installed
if [[ ! -f /mnt/boot/vmlinuz-$KERNEL_NAME ]]; then
    error "Kernel vmlinuz-$KERNEL_NAME not found in /mnt/boot/"
fi
log "Found kernel: vmlinuz-$KERNEL_NAME"

# Configure Limine bootloader using Omarchy's approach
log "Configuring Limine bootloader..."
mkdir -p /mnt/boot/EFI/BOOT /mnt/boot/EFI/limine

# Copy Limine EFI files
cp /mnt/usr/share/limine/BOOTX64.EFI /mnt/boot/EFI/BOOT/
cp /mnt/usr/share/limine/BOOTX64.EFI /mnt/boot/EFI/limine/

# Create /etc/default/limine for limine-update (Omarchy style)
CRYPT_CMDLINE="cryptdevice=UUID=$LUKS_UUID:cryptroot root=/dev/mapper/cryptroot rootflags=subvol=@"
cat > /mnt/etc/default/limine << LIMINEDEF
TARGET_OS_NAME="Omarchy"

ESP_PATH="/boot"

KERNEL_CMDLINE[default]="$CRYPT_CMDLINE rw quiet splash"

ENABLE_UKI=yes
CUSTOM_UKI_NAME="omarchy"

ENABLE_LIMINE_FALLBACK=yes

# Find and add other bootloaders (Windows)
FIND_BOOTLOADERS=yes

BOOT_ORDER="*, *fallback, Snapshots"

MAX_SNAPSHOT_ENTRIES=5

SNAPSHOT_FORMAT_CHOICE=5
LIMINEDEF

# Create base limine.conf (limine-update will add entries)
cat > /mnt/boot/limine.conf << 'LIMINE'
default_entry: 2
interface_branding: Omarchy Bootloader
interface_branding_color: 2
hash_mismatch_panic: no

term_background: 1a1b26
backdrop: 1a1b26

term_palette: 15161e;f7768e;9ece6a;e0af68;7aa2f7;bb9af7;7dcfff;a9b1d6
term_palette_bright: 414868;f7768e;9ece6a;e0af68;7aa2f7;bb9af7;7dcfff;c0caf5

term_foreground: c0caf5
term_foreground_bright: c0caf5
term_background_bright: 24283b
LIMINE

# Run limine-update to generate boot entries (including snapshots)
log "Running limine-update to generate boot entries..."
arch-chroot /mnt limine-update || warn "limine-update had issues"

# Enable btrfs quota for space-aware snapshot cleanup
log "Enabling btrfs quota..."
btrfs quota enable /mnt || warn "Could not enable btrfs quota"

#===============================================================================
# INSTALL OMARCHY
#===============================================================================

header "INSTALLING OMARCHY PACKAGES"

# Mount offline mirror if available
if [[ -d /var/cache/omarchy/mirror/offline ]]; then
    mkdir -p /mnt/var/cache/omarchy/mirror/offline
    mount --bind /var/cache/omarchy/mirror/offline /mnt/var/cache/omarchy/mirror/offline
fi

if [[ -d /opt/packages ]]; then
    mkdir -p /mnt/opt/packages
    mount --bind /opt/packages /mnt/opt/packages
fi

# Backup system pacman.conf and use ISO's version temporarily for offline install
cp /mnt/etc/pacman.conf /mnt/etc/pacman.conf.bak
cp /etc/pacman.conf /mnt/etc/pacman.conf

# Copy omarchy to user's home
mkdir -p /mnt/home/$OMARCHY_USER/.local/share/
cp -a /root/omarchy /mnt/home/$OMARCHY_USER/.local/share/
# Ensure all scripts are executable
chmod -R +x /mnt/home/$OMARCHY_USER/.local/share/omarchy/bin/
chmod -R +x /mnt/home/$OMARCHY_USER/.local/share/omarchy/install/
find /mnt/home/$OMARCHY_USER/.local/share/omarchy -name "*.sh" -exec chmod +x {} \;
arch-chroot /mnt chown -R $OMARCHY_USER:$OMARCHY_USER /home/$OMARCHY_USER/.local/

# Get user info from configurator-generated files
USER_FULL_NAME=""
USER_EMAIL=""
[[ -f /root/user_full_name.txt ]] && USER_FULL_NAME=$(<user_full_name.txt)
[[ -f /root/user_email_address.txt ]] && USER_EMAIL=$(<user_email_address.txt)
OMARCHY_MIRROR=""
[[ -f /root/omarchy_mirror ]] && OMARCHY_MIRROR=$(<omarchy_mirror)

# Run omarchy installer in chroot (with proper environment like original installer)
log "Installing Omarchy packages and configuration..."
arch-chroot /mnt /bin/bash -c "
    export HOME=/home/$OMARCHY_USER
    export USER=$OMARCHY_USER
    export OMARCHY_CHROOT_INSTALL=1
    export OMARCHY_USER_NAME='$USER_FULL_NAME'
    export OMARCHY_USER_EMAIL='$USER_EMAIL'
    export OMARCHY_MIRROR='$OMARCHY_MIRROR'
    cd /home/$OMARCHY_USER/.local/share/omarchy
    sudo -u $OMARCHY_USER -E bash -c 'source install.sh' || true
" || warn "Omarchy install script had some issues (this may be OK)"

# Set up proper pacman.conf for the installed system
log "Configuring pacman for installed system..."
cat > /mnt/etc/pacman.conf << 'PACMANCONF'
[options]
Color
ILoveCandy
VerbosePkgLists
HoldPkg = pacman glibc
Architecture = auto
CheckSpace
ParallelDownloads = 5
DownloadUser = alpm

SigLevel = Required DatabaseOptional
LocalFileSigLevel = Optional

[core]
Include = /etc/pacman.d/mirrorlist

[extra]
Include = /etc/pacman.d/mirrorlist

[multilib]
Include = /etc/pacman.d/mirrorlist

[omarchy]
SigLevel = Optional TrustAll
Server = https://pkgs.omarchy.org/stable/$arch
PACMANCONF

# Set up mirrorlist
cat > /mnt/etc/pacman.d/mirrorlist << 'MIRRORLIST'
Server = https://stable-mirror.omarchy.org/$repo/os/$arch
MIRRORLIST

log "Pacman configured with Omarchy stable repos"

# Now that Omarchy is installed, set Plymouth theme and rebuild initramfs
log "Setting Omarchy Plymouth theme..."
if [[ -d /mnt/usr/share/plymouth/themes/omarchy ]]; then
    arch-chroot /mnt plymouth-set-default-theme omarchy
    log "Rebuilding initramfs with Omarchy Plymouth theme..."
    arch-chroot /mnt mkinitcpio -P
else
    warn "Omarchy Plymouth theme not found"
fi

#===============================================================================
# FINAL SETUP
#===============================================================================

header "FINAL CONFIGURATION"

# Enable services
log "Enabling services..."
arch-chroot /mnt systemctl enable sddm 2>/dev/null || true
arch-chroot /mnt systemctl enable NetworkManager 2>/dev/null || true
arch-chroot /mnt systemctl enable iwd 2>/dev/null || true
arch-chroot /mnt systemctl enable bluetooth 2>/dev/null || true

# Snapper configuration
log "Installing and configuring snapper..."
arch-chroot /mnt pacman --noconfirm --needed -S snapper 2>/dev/null || true

# Create snapper config for root - this creates .snapshots as nested subvolume in @
log "Creating snapper root config..."
arch-chroot /mnt snapper -c root create-config / || warn "snapper create-config had issues"

# Enable snapper services
arch-chroot /mnt systemctl enable snapper-cleanup.timer 2>/dev/null || true
arch-chroot /mnt systemctl enable snapper-timeline.timer 2>/dev/null || true
arch-chroot /mnt systemctl enable limine-snapper-sync.service 2>/dev/null || true

# Create initial snapshot
log "Creating initial snapshot..."
arch-chroot /mnt snapper -c root create -d "Fresh install" || warn "Could not create initial snapshot"

# Sync snapshots to limine boot menu
log "Syncing snapshots to limine..."
arch-chroot /mnt limine-snapper-sync 2>/dev/null || warn "limine-snapper-sync had issues (may need to run after first boot)"

log "Snapper configured with bootable snapshots"

# Configure SDDM auto-login (Omarchy uses hyprland-uwsm session)
log "Configuring auto-login for $OMARCHY_USER..."
mkdir -p /mnt/etc/sddm.conf.d
cat > /mnt/etc/sddm.conf.d/autologin.conf << SDDMCONF
[Autologin]
User=$OMARCHY_USER
Session=hyprland-uwsm
[Theme]
Current=breeze
SDDMCONF

# Create UEFI boot entry
log "Creating UEFI boot entry..."
DRIVE=$(cat /tmp/install_drive)
EFI_PART=$(cat /tmp/efi_part)
EFI_PART_NUM=$(echo "$EFI_PART" | grep -oE '[0-9]+$')

# Remove any existing Omarchy or Limine entries to avoid duplicates
for entry in $(efibootmgr 2>/dev/null | grep -iE "Omarchy|Limine" | sed -n 's/^Boot\([0-9A-Fa-f]*\).*/\1/p'); do
    log "Removing old boot entry: $entry"
    efibootmgr -b "$entry" -B 2>/dev/null || true
done

# Create single Omarchy entry
efibootmgr --create --disk "$DRIVE" --part "$EFI_PART_NUM" \
    --loader '\\EFI\\BOOT\\BOOTX64.EFI' --label 'Omarchy' 2>/dev/null || \
    warn "Could not create UEFI entry - use F12 boot menu"

#===============================================================================
# CLEANUP
#===============================================================================

header "CLEANING UP"

umount /mnt/var/cache/omarchy/mirror/offline 2>/dev/null || true
umount /mnt/opt/packages 2>/dev/null || true
umount -R /mnt
cryptsetup close cryptroot

rm -f /tmp/luks_uuid /tmp/efi_part /tmp/root_part /tmp/install_drive /tmp/kernel_name

#===============================================================================
# DONE
#===============================================================================

header "INSTALLATION COMPLETE!"

echo ""
echo -e "${GREEN}╔══════════════════════════════════════════════════════════════╗${NC}"
echo -e "${GREEN}║                                                              ║${NC}"
echo -e "${GREEN}║   Omarchy has been installed alongside Windows!              ║${NC}"
echo -e "${GREEN}║                                                              ║${NC}"
echo -e "${GREEN}╚══════════════════════════════════════════════════════════════╝${NC}"
echo ""
echo "To boot:"
echo "  1. Reboot your computer"
echo "  2. Press F12/F8/DEL at startup for boot menu"
echo "  3. Select 'Omarchy' for Linux"
echo "  4. Select 'Windows Boot Manager' for Windows"
echo ""
echo "  Username: $OMARCHY_USER"
echo ""
read -p "Press Enter to reboot..." _
reboot
DUALBOOT_SCRIPT

chmod +x "$WORK_DIR/extracted/root/dualboot-setup.sh"

#===============================================================================
# MODIFY BOOT SEQUENCE
#===============================================================================

log "Configuring boot sequence..."

# Backup original automated script
if [[ -f "$WORK_DIR/extracted/root/.automated_script.sh" ]]; then
    mv "$WORK_DIR/extracted/root/.automated_script.sh" "$WORK_DIR/extracted/root/.automated_script.sh.orig"
fi

# Replace .zlogin to run our dual-boot setup
cat > "$WORK_DIR/extracted/root/.zlogin" << 'EOF'
# Omarchy Dual-Boot Installer
if [[ $(tty) == "/dev/tty1" ]]; then
    /root/dualboot-setup.sh
fi
EOF

log "Boot sequence configured"

#===============================================================================
# REPACK SQUASHFS
#===============================================================================

log "Repacking squashfs (this takes several minutes)..."
rm "$SQUASHFS"
mksquashfs "$WORK_DIR/extracted" "$SQUASHFS" -comp zstd -Xcompression-level 3

if [[ ! -f "$SQUASHFS" ]]; then
    error "Failed to create squashfs!"
fi
log "New squashfs created: $(du -h "$SQUASHFS" | cut -f1)"

#===============================================================================
# UPDATE BOOT CONFIGS
#===============================================================================

log "Creating new ISO..."

ISO_LABEL=$(isoinfo -d -i "$SOURCE_ISO" 2>/dev/null | grep "Volume id:" | cut -d: -f2 | tr -d ' ' || echo "OMARCHY")
log "ISO label: $ISO_LABEL"

# Update boot configs to use label instead of UUID
log "Updating boot configurations..."
for cfg in "$WORK_DIR/newiso/boot/syslinux/"*.cfg \
           "$WORK_DIR/newiso/loader/entries/"*.conf \
           "$WORK_DIR/newiso/boot/grub/grub.cfg" \
           "$WORK_DIR/newiso/EFI/BOOT/grub.cfg"; do
    if [[ -f "$cfg" ]]; then
        sed -i "s/archisosearchuuid=[^ ]*/archisolabel=$ISO_LABEL/g" "$cfg"
    fi
done

#===============================================================================
# EXTRACT EFI BOOT IMAGE FROM ORIGINAL ISO
#===============================================================================

log "Extracting EFI boot image from original ISO..."

# Get the xorriso parameters from the original ISO to replicate its boot setup
XORRISO_PARAMS=$(xorriso -indev "$SOURCE_ISO" -report_el_torito as_mkisofs 2>/dev/null)

# Extract the appended EFI partition from the original ISO
# Format: -append_partition 2 0xef --interval:local_fs:STARTd-ENDd::
EFI_LINE=$(echo "$XORRISO_PARAMS" | grep "append_partition 2 0xef")
if [[ -n "$EFI_LINE" ]]; then
    log "Found EFI partition in original ISO"
    # Extract start and end from the interval (format: NUMBERd-NUMBERd)
    EFI_RANGE=$(echo "$EFI_LINE" | grep -oE '[0-9]+d-[0-9]+d')
    EFI_START=$(echo "$EFI_RANGE" | cut -d'-' -f1 | tr -d 'd')
    EFI_END=$(echo "$EFI_RANGE" | cut -d'-' -f2 | tr -d 'd')

    if [[ -n "$EFI_START" && -n "$EFI_END" ]]; then
        EFI_SIZE=$((EFI_END - EFI_START + 1))
        log "Extracting EFI partition: start=$EFI_START end=$EFI_END size=$EFI_SIZE sectors"
        # The 'd' values from xorriso are in 512-byte sectors
        dd if="$SOURCE_ISO" of="$WORK_DIR/efi.img" bs=512 skip="$EFI_START" count="$EFI_SIZE" status=progress

        if [[ -f "$WORK_DIR/efi.img" ]]; then
            EFI_IMG_SIZE=$(stat -c%s "$WORK_DIR/efi.img")
            log "EFI image created: $((EFI_IMG_SIZE / 1024 / 1024))MB"
        else
            warn "Failed to extract EFI image"
        fi
    else
        warn "Could not parse EFI partition range"
    fi
else
    warn "No EFI partition found in original ISO"
fi

#===============================================================================
# CREATE ISO
#===============================================================================

if [[ -f "$WORK_DIR/newiso/boot/syslinux/isolinux.bin" ]]; then
    log "Creating hybrid BIOS/UEFI ISO..."

    if [[ -f "$WORK_DIR/efi.img" ]]; then
        # Use extracted EFI image for proper UEFI boot
        xorriso -as mkisofs \
            -iso-level 3 \
            -full-iso9660-filenames \
            -rational-rock \
            -volid "$ISO_LABEL" \
            -eltorito-boot boot/syslinux/isolinux.bin \
            -eltorito-catalog boot/syslinux/boot.cat \
            -no-emul-boot \
            -boot-load-size 4 \
            -boot-info-table \
            -isohybrid-mbr "$WORK_DIR/newiso/boot/syslinux/isohdpfx.bin" \
            -eltorito-alt-boot \
            -e --interval:appended_partition_2:all:: \
            -no-emul-boot \
            -append_partition 2 0xef "$WORK_DIR/efi.img" \
            -isohybrid-gpt-basdat \
            -o "$OUTPUT_ISO" \
            "$WORK_DIR/newiso"
    else
        warn "EFI image not found, using fallback method..."
        xorriso -as mkisofs \
            -iso-level 3 \
            -full-iso9660-filenames \
            -rational-rock \
            -volid "$ISO_LABEL" \
            -eltorito-boot boot/syslinux/isolinux.bin \
            -eltorito-catalog boot/syslinux/boot.cat \
            -no-emul-boot \
            -boot-load-size 4 \
            -boot-info-table \
            -isohybrid-mbr "$WORK_DIR/newiso/boot/syslinux/isohdpfx.bin" \
            -eltorito-alt-boot \
            -e EFI/BOOT/BOOTX64.EFI \
            -no-emul-boot \
            -isohybrid-gpt-basdat \
            -o "$OUTPUT_ISO" \
            "$WORK_DIR/newiso"
    fi
else
    log "Creating UEFI-only ISO..."
    xorriso -as mkisofs \
        -iso-level 3 \
        -full-iso9660-filenames \
        -volid "$ISO_LABEL" \
        -e EFI/BOOT/BOOTX64.EFI \
        -no-emul-boot \
        -isohybrid-gpt-basdat \
        -o "$OUTPUT_ISO" \
        "$WORK_DIR/newiso"
fi

#===============================================================================
# DONE
#===============================================================================

echo ""
echo -e "${GREEN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo -e "${GREEN}  SUCCESS!${NC}"
echo -e "${GREEN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo ""
echo "  Created: $OUTPUT_ISO"
echo "  Size:    $(du -h "$OUTPUT_ISO" | cut -f1)"
echo ""
echo "  This ISO provides:"
echo "    • Option 1: Dual-boot install (preserves Windows)"
echo "    • Option 3: Standard install (original Omarchy installer)"
echo ""
echo "  Next steps:"
echo "    1. Copy ISO to Ventoy USB drive"
echo "    2. Boot from Ventoy"
echo "    3. Select option 1 for dual-boot"
echo ""