#!/bin/bash
# Create Truly Bootable Backup USB with Alpine Linux
# This creates a complete bootable environment that preserves existing boot capability

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check if running as root
if [[ $EUID -eq 0 ]]; then
    print_error "Do not run as root. Script will use sudo when needed."
    exit 1
fi

print_status "Bootable Backup USB Creator (Preserves Existing Boot)"
echo "=========================================================="
echo

# Check current USB drive status
USB_DRIVE="/dev/sda"  # Your current USB

print_status "Analyzing current USB structure..."
lsblk "$USB_DRIVE"
echo

# Check if it already has a bootable system
HAS_BOOT=$(lsblk "$USB_DRIVE" | grep -i boot || true)
if [[ -n "$HAS_BOOT" ]]; then
    print_warning "USB appears to have existing boot partition"
    print_warning "We'll preserve this and add backup functionality"
else
    print_status "No existing boot detected - will create complete bootable system"
fi

read -p "Continue to make this USB fully bootable with backup functionality? (yes/no): " confirm
if [[ "$confirm" != "yes" ]]; then
    exit 0
fi

# Create temporary directory for Alpine Linux
WORK_DIR="/tmp/alpine_usb_$$"
mkdir -p "$WORK_DIR"
cd "$WORK_DIR"

print_status "Downloading Alpine Linux (minimal, fast-booting)..."

# Download Alpine Linux (very small, perfect for this)
ALPINE_VERSION="3.18.4"
ALPINE_ISO="alpine-standard-${ALPINE_VERSION}-x86_64.iso"

if [[ ! -f "$ALPINE_ISO" ]]; then
    wget "https://dl-cdn.alpinelinux.org/alpine/v3.18/releases/x86_64/$ALPINE_ISO" || {
        print_error "Failed to download Alpine Linux"
        exit 1
    }
fi

print_status "Preparing bootable USB with backup functionality..."

# Mount the ISO to extract files
ISO_MOUNT="$WORK_DIR/iso_mount"
mkdir -p "$ISO_MOUNT"
sudo mount -o loop "$ALPINE_ISO" "$ISO_MOUNT"

# Create mount points for USB partitions
BOOT_MOUNT="$WORK_DIR/usb_boot"
DATA_MOUNT="$WORK_DIR/usb_data"
mkdir -p "$BOOT_MOUNT" "$DATA_MOUNT"

# Mount USB partitions (created by previous script)
sudo mount "${USB_DRIVE}1" "$BOOT_MOUNT"
sudo mount "${USB_DRIVE}2" "$DATA_MOUNT"

print_status "Installing Alpine Linux boot files..."

# Copy Alpine Linux files to boot partition
sudo cp -r "$ISO_MOUNT"/* "$BOOT_MOUNT/"

# Install GRUB for UEFI/BIOS boot
print_status "Installing GRUB bootloader..."
sudo grub-install --target=x86_64-efi --efi-directory="$BOOT_MOUNT" --boot-directory="$BOOT_MOUNT/boot" --removable --force

# Create GRUB configuration for automatic backup boot
sudo tee "$BOOT_MOUNT/boot/grub/grub.cfg" > /dev/null << 'EOF'
set timeout=10
set default=0

menuentry "Automatic System Backup" {
    linux /boot/vmlinuz-lts modules=loop,squashfs,sd-mod,usb-storage quiet alpine_dev=sda2:/backup-tools/alpine.apkovl.tar.gz
    initrd /boot/initramfs-lts
}

menuentry "Manual Backup Mode" {
    linux /boot/vmlinuz-lts modules=loop,squashfs,sd-mod,usb-storage alpine_dev=sda2:/backup-tools/alpine.apkovl.tar.gz
    initrd /boot/initramfs-lts
}

menuentry "Alpine Linux (Standard)" {
    linux /boot/vmlinuz-lts modules=loop,squashfs,sd-mod,usb-storage
    initrd /boot/initramfs-lts
}
EOF

print_status "Creating backup environment overlay..."

# Create Alpine overlay for automatic backup
OVERLAY_DIR="$WORK_DIR/overlay"
mkdir -p "$OVERLAY_DIR/etc/init.d"
mkdir -p "$OVERLAY_DIR/usr/local/bin"
mkdir -p "$OVERLAY_DIR/etc/runlevels/default"

# Create backup script that runs on boot
sudo tee "$OVERLAY_DIR/usr/local/bin/backup-menu" > /dev/null << 'EOF'
#!/bin/sh
# Interactive backup menu

clear
echo "========================================"
echo "      AUTOMATIC BACKUP SYSTEM"
echo "========================================"
echo
echo "Available drives:"
lsblk -d -o NAME,SIZE,TYPE,TRAN | grep -E "(disk|NAME)"
echo

# Auto-detect drives
INTERNAL_CANDIDATES=$(lsblk -d -n -o NAME,TYPE,TRAN | grep "disk" | grep -v "usb" | awk '{print "/dev/" $1}')
EXTERNAL_CANDIDATES=$(lsblk -d -n -o NAME,TYPE,TRAN | grep "disk" | grep "usb" | awk '{print "/dev/" $1}' | grep -v sda)

echo "Drive Selection:"
echo "=================="

# Select source drive
echo "Available internal drives:"
echo "$INTERNAL_CANDIDATES" | nl -v 1
echo
read -p "Select source drive number (or enter path): " SOURCE_CHOICE

if echo "$SOURCE_CHOICE" | grep -q "^[0-9]"; then
    SOURCE_DRIVE=$(echo "$INTERNAL_CANDIDATES" | sed -n "${SOURCE_CHOICE}p")
else
    SOURCE_DRIVE="$SOURCE_CHOICE"
fi

# Select target drive
echo
echo "Available external drives:"
echo "$EXTERNAL_CANDIDATES" | nl -v 1
echo
read -p "Select target drive number (or enter path): " TARGET_CHOICE

if echo "$TARGET_CHOICE" | grep -q "^[0-9]"; then
    TARGET_DRIVE=$(echo "$EXTERNAL_CANDIDATES" | sed -n "${TARGET_CHOICE}p")
else
    TARGET_DRIVE="$TARGET_CHOICE"
fi

echo
echo "BACKUP CONFIGURATION:"
echo "===================="
echo "Source (will be copied FROM): $SOURCE_DRIVE"
echo "Target (will be overwritten): $TARGET_DRIVE"
echo
echo "⚠️  ALL DATA ON $TARGET_DRIVE WILL BE DESTROYED! ⚠️"
echo
read -p "Continue with backup? (yes/no): " CONFIRM

if [ "$CONFIRM" != "yes" ]; then
    echo "Backup cancelled"
    read -p "Press Enter to shutdown..."
    poweroff
    exit 0
fi

echo
echo "Starting backup..."
echo "=================="

# Get drive sizes for progress
SOURCE_SIZE=$(blockdev --getsize64 "$SOURCE_DRIVE")
SOURCE_SIZE_GB=$((SOURCE_SIZE / 1024 / 1024 / 1024))

echo "Copying $SOURCE_SIZE_GB GB from $SOURCE_DRIVE to $TARGET_DRIVE"
echo

# Perform backup with progress
if command -v pv >/dev/null 2>&1; then
    dd if="$SOURCE_DRIVE" bs=4M | pv -s "$SOURCE_SIZE" | dd of="$TARGET_DRIVE" bs=4M conv=fdatasync
else
    dd if="$SOURCE_DRIVE" of="$TARGET_DRIVE" bs=4M status=progress conv=fdatasync
fi

if [ $? -eq 0 ]; then
    echo
    echo "========================================"
    echo "    BACKUP COMPLETED SUCCESSFULLY!"
    echo "========================================"
    echo "Source: $SOURCE_DRIVE"
    echo "Target: $TARGET_DRIVE"
    echo "Size: $SOURCE_SIZE_GB GB"
else
    echo
    echo "========================================"
    echo "        BACKUP FAILED!"
    echo "========================================"
fi

echo
echo "System will shutdown in 30 seconds..."
echo "Press Ctrl+C to cancel shutdown"
sleep 30
poweroff
EOF

chmod +x "$OVERLAY_DIR/usr/local/bin/backup-menu"

# Create init script to run backup on boot
sudo tee "$OVERLAY_DIR/etc/init.d/autobackup" > /dev/null << 'EOF'
#!/sbin/openrc-run

name="autobackup"
description="Automatic backup system"

depend() {
    need localmount
    after *
}

start() {
    ebegin "Starting automatic backup system"
    
    # Wait for devices to settle
    sleep 5
    
    # Check kernel command line for auto mode
    if grep -q "autobackup" /proc/cmdline; then
        /usr/local/bin/backup-menu
    else
        # Manual mode - provide choice
        echo "Press 'b' for backup or any other key to continue to shell..."
        read -t 10 -n 1 choice
        if [ "$choice" = "b" ] || [ "$choice" = "B" ]; then
            /usr/local/bin/backup-menu
        fi
    fi
    
    eend $?
}
EOF

chmod +x "$OVERLAY_DIR/etc/init.d/autobackup"

# Enable the service
ln -sf /etc/init.d/autobackup "$OVERLAY_DIR/etc/runlevels/default/autobackup"

# Create the overlay archive
cd "$OVERLAY_DIR"
tar czf "$DATA_MOUNT/alpine.apkovl.tar.gz" *

# Copy our backup tools to the data partition as well
SCRIPT_DIR=$(dirname "$(realpath "$0")")
sudo cp "$SCRIPT_DIR"/*.sh "$DATA_MOUNT/"

print_status "Installing additional tools..."

# Create an Alpine package cache with useful tools
sudo mkdir -p "$DATA_MOUNT/apk-cache"

# Cleanup
cd "$WORK_DIR/.."
sudo umount "$ISO_MOUNT" "$BOOT_MOUNT" "$DATA_MOUNT" 2>/dev/null || true
rm -rf "$WORK_DIR"

print_success "Fully bootable backup USB created!"
print_success "USB: $USB_DRIVE"
echo
print_success "WHAT YOU CAN NOW DO:"
print_success "1. Boot from this USB in BIOS/UEFI menu"
print_success "2. Select 'Automatic System Backup' from GRUB menu"
print_success "3. Choose source and target drives from interactive menu"
print_success "4. Backup runs automatically with progress display"
print_success "5. System shuts down when complete"
echo
print_warning "The USB remains bootable for future backups!"
print_warning "Your backup tools are preserved in the data partition"
EOF