#!/bin/bash
# Create Clonezilla-based Backup USB
# This adapts Clonezilla Live with our custom backup functionality

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"
}

print_status "Clonezilla-based Backup USB Creator"
echo "======================================"
echo

print_warning "This will enhance Clonezilla Live with our custom backup tools"
print_warning "You'll get both Clonezilla functionality AND our automated backup"
echo

# Check if we have a USB drive
USB_DRIVE="/dev/sda"
if [[ ! -b "$USB_DRIVE" ]]; then
    print_error "USB drive $USB_DRIVE not found!"
    exit 1
fi

print_status "Current USB: $USB_DRIVE"
lsblk "$USB_DRIVE"
echo

# Download Clonezilla if not already present
CLONEZILLA_ISO="clonezilla-live-3.1.0-22-amd64.iso"
# Using OSDN mirror which is typically faster than SourceForge
CLONEZILLA_URL="https://osdn.net/dl/clonezilla/$CLONEZILLA_ISO"

if [[ ! -f "$CLONEZILLA_ISO" ]]; then
    print_status "Downloading Clonezilla Live from OSDN mirror..."
    wget "$CLONEZILLA_URL" || {
        print_warning "OSDN mirror failed, trying GitHub mirror..."
        CLONEZILLA_URL="https://github.com/stevenshiau/clonezilla/releases/download/3.1.0-22/$CLONEZILLA_ISO"
        wget "$CLONEZILLA_URL" || {
            print_warning "GitHub mirror failed, trying SourceForge as fallback..."
            CLONEZILLA_URL="https://downloads.sourceforge.net/clonezilla/$CLONEZILLA_ISO"
            wget "$CLONEZILLA_URL" || {
                print_error "Failed to download Clonezilla from all mirrors"
                exit 1
            }
        }
    }
fi

print_warning "⚠️  This will COMPLETELY RECREATE the USB drive with Clonezilla! ⚠️"
print_warning "All current data will be lost!"
echo
read -p "Continue? (yes/no): " confirm

if [[ "$confirm" != "yes" ]]; then
    print_error "Operation cancelled"
    exit 1
fi

print_status "Creating Clonezilla Live USB with backup tools..."

# Unmount any mounted partitions
sudo umount "${USB_DRIVE}"* 2>/dev/null || true

# Create new partition table
print_status "Creating partition table..."
sudo parted "$USB_DRIVE" --script mklabel msdos
sudo parted "$USB_DRIVE" --script mkpart primary fat32 1MiB 4GiB
sudo parted "$USB_DRIVE" --script mkpart primary ext4 4GiB 100%
sudo parted "$USB_DRIVE" --script set 1 boot on

# Format partitions
print_status "Formatting partitions..."
if [[ "$USB_DRIVE" == *"nvme"* ]]; then
    BOOT_PART="${USB_DRIVE}p1"
    DATA_PART="${USB_DRIVE}p2"
else
    BOOT_PART="${USB_DRIVE}1"
    DATA_PART="${USB_DRIVE}2"
fi

sudo mkfs.fat -F32 -n "CLONEZILLA" "$BOOT_PART"
sudo mkfs.ext4 -L "BACKUP_TOOLS" "$DATA_PART"

# Mount partitions
BOOT_MOUNT="/tmp/clonezilla_boot_$$"
DATA_MOUNT="/tmp/clonezilla_data_$$"
ISO_MOUNT="/tmp/clonezilla_iso_$$"

sudo mkdir -p "$BOOT_MOUNT" "$DATA_MOUNT" "$ISO_MOUNT"
sudo mount "$BOOT_PART" "$BOOT_MOUNT"
sudo mount "$DATA_PART" "$DATA_MOUNT"
sudo mount -o loop "$CLONEZILLA_ISO" "$ISO_MOUNT"

print_status "Installing Clonezilla Live..."

# Copy Clonezilla files
sudo cp -r "$ISO_MOUNT"/* "$BOOT_MOUNT/"

# Install GRUB
print_status "Installing GRUB bootloader..."
sudo grub-install --target=i386-pc --boot-directory="$BOOT_MOUNT/boot" "$USB_DRIVE"

# Create enhanced GRUB configuration
print_status "Creating enhanced GRUB menu..."
sudo tee "$BOOT_MOUNT/boot/grub/grub.cfg" > /dev/null << 'EOF'
set timeout=15
set default=0

menuentry "🚀 AUTOMATIC SYSTEM BACKUP" {
    linux /live/vmlinuz boot=live union=overlay username=user config components quiet noswap edd=on nomodeset ocs_live_run="ocs-live-general" ocs_live_extra_param="" keyboard-layouts= ocs_live_batch="no" locales= vga=788 ip= net.ifnames=0 nosplash i915.blacklist=yes radeonhd.blacklist=yes nouveau.blacklist=yes vmwgfx.enable_fbdev=1 systemd.show_status=0
    initrd /live/initrd.img
}

menuentry "🔧 MANUAL BACKUP MODE" {
    linux /live/vmlinuz boot=live union=overlay username=user config components quiet noswap edd=on nomodeset ocs_live_run="ocs-live-general" ocs_live_extra_param="" keyboard-layouts= ocs_live_batch="no" locales= vga=788 ip= net.ifnames=0 nosplash i915.blacklist=yes radeonhd.blacklist=yes nouveau.blacklist=yes vmwgfx.enable_fbdev=1 systemd.show_status=0 custom_backup=manual
    initrd /live/initrd.img
}

menuentry "📦 CLONEZILLA LIVE (Original)" {
    linux /live/vmlinuz boot=live union=overlay username=user config components quiet noswap edd=on nomodeset ocs_live_run="ocs-live-general" ocs_live_extra_param="" keyboard-layouts= ocs_live_batch="no" locales= vga=788 ip= net.ifnames=0 nosplash i915.blacklist=yes radeonhd.blacklist=yes nouveau.blacklist=yes vmwgfx.enable_fbdev=1
    initrd /live/initrd.img
}

menuentry "🛠️  CLONEZILLA EXPERT MODE" {
    linux /live/vmlinuz boot=live union=overlay username=user config components quiet noswap edd=on nomodeset ocs_live_run="ocs-expert" ocs_live_extra_param="" keyboard-layouts= ocs_live_batch="no" locales= vga=788 ip= net.ifnames=0 nosplash i915.blacklist=yes radeonhd.blacklist=yes nouveau.blacklist=yes vmwgfx.enable_fbdev=1
    initrd /live/initrd.img
}

menuentry "🔍 MEMORY TEST" {
    linux16 /live/memtest
}
EOF

print_status "Installing backup tools to data partition..."

# Copy our backup scripts
SCRIPT_DIR=$(dirname "$(realpath "$0")")
sudo cp "$SCRIPT_DIR"/*.sh "$DATA_MOUNT/"
sudo chmod +x "$DATA_MOUNT"/*.sh

# Also place the automated script in the Clonezilla filesystem for direct access
sudo mkdir -p "$BOOT_MOUNT/live/image/backup_tools"
sudo cp "$SCRIPT_DIR/automated_clonezilla_backup.sh" "$BOOT_MOUNT/live/image/backup_tools/"
sudo chmod +x "$BOOT_MOUNT/live/image/backup_tools/automated_clonezilla_backup.sh"

# Create custom startup script that checks for auto_backup parameter
sudo tee "$BOOT_MOUNT/live/image/auto-start.sh" > /dev/null << 'EOF'
#!/bin/bash
# Check if auto_backup=true is in kernel parameters
if grep -q "auto_backup=true" /proc/cmdline; then
    echo "Starting automatic backup mode..."
    sleep 2
    # Mount backup partition and run our script
    mkdir -p /tmp/backup_mount
    mount /dev/sda2 /tmp/backup_mount 2>/dev/null
    if [ -f /tmp/backup_mount/automated_clonezilla_backup.sh ]; then
        /tmp/backup_mount/automated_clonezilla_backup.sh
    else
        echo "Backup script not found, starting normal Clonezilla"
        ocs-live-general
    fi
else
    # Normal Clonezilla startup
    ocs-live-general
fi
EOF

sudo chmod +x "$BOOT_MOUNT/live/image/auto-start.sh"

# Create startup script for Clonezilla
sudo tee "$DATA_MOUNT/auto-backup.sh" > /dev/null << 'EOF'
#!/bin/bash
# Auto-backup script for Clonezilla Live

export PATH=/usr/bin:/bin:/sbin:/usr/sbin:/usr/local/bin

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

clear
echo -e "${GREEN}"
echo "================================================"
echo "    CLONEZILLA-BASED AUTOMATIC BACKUP"
echo "================================================"
echo -e "${NC}"

# Check if we're in auto mode
AUTO_MODE="false"
if grep -q "custom_backup=auto" /proc/cmdline; then
    AUTO_MODE="true"
fi

# Wait for drives to settle
echo "Waiting for drives to be detected..."
sleep 5

# Function to list drives
list_drives() {
    echo -e "${BLUE}Available drives:${NC}"
    echo "=================="
    lsblk -d -o NAME,SIZE,TYPE,TRAN,MODEL | grep -E "(NAME|disk)"
    echo
}

# Function to get drive selection
select_drives() {
    list_drives
    
    echo -e "${YELLOW}SELECT SOURCE DRIVE (internal, to backup FROM):${NC}"
    INTERNAL_DRIVES=$(lsblk -d -n -o NAME,TRAN | grep -v usb | grep -v loop | awk '{print $1}')
    echo "$INTERNAL_DRIVES" | nl -v 1
    echo
    read -p "Enter source drive number or path: " SOURCE_INPUT
    
    if echo "$SOURCE_INPUT" | grep -q "^[0-9]"; then
        SOURCE_DRIVE="/dev/$(echo "$INTERNAL_DRIVES" | sed -n "${SOURCE_INPUT}p")"
    else
        SOURCE_DRIVE="$SOURCE_INPUT"
    fi
    
    echo
    echo -e "${YELLOW}SELECT TARGET DRIVE (external, will be OVERWRITTEN):${NC}"
    EXTERNAL_DRIVES=$(lsblk -d -n -o NAME,TRAN | grep -v loop | awk '{print $1}')
    echo "$EXTERNAL_DRIVES" | nl -v 1
    echo
    read -p "Enter target drive number or path: " TARGET_INPUT
    
    if echo "$TARGET_INPUT" | grep -q "^[0-9]"; then
        TARGET_DRIVE="/dev/$(echo "$EXTERNAL_DRIVES" | sed -n "${TARGET_INPUT}p")"
    else
        TARGET_DRIVE="$TARGET_INPUT"
    fi
}

# Function to perform backup
perform_backup() {
    echo
    echo -e "${GREEN}BACKUP CONFIGURATION:${NC}"
    echo "===================="
    echo "Source: $SOURCE_DRIVE"
    echo "Target: $TARGET_DRIVE"
    echo
    
    # Show sizes
    if [[ -b "$SOURCE_DRIVE" ]]; then
        SOURCE_SIZE=$(blockdev --getsize64 "$SOURCE_DRIVE" 2>/dev/null || echo "0")
        SOURCE_GB=$((SOURCE_SIZE / 1024 / 1024 / 1024))
        echo "Source size: ${SOURCE_GB}GB"
    fi
    
    if [[ -b "$TARGET_DRIVE" ]]; then
        TARGET_SIZE=$(blockdev --getsize64 "$TARGET_DRIVE" 2>/dev/null || echo "0")
        TARGET_GB=$((TARGET_SIZE / 1024 / 1024 / 1024))
        echo "Target size: ${TARGET_GB}GB"
    fi
    
    echo
    echo -e "${RED}⚠️  ALL DATA ON $TARGET_DRIVE WILL BE DESTROYED! ⚠️${NC}"
    echo
    
    if [[ "$AUTO_MODE" == "true" ]]; then
        echo "Auto mode - starting backup in 10 seconds..."
        echo "(Press Ctrl+C to cancel)"
        sleep 10
    else
        read -p "Continue with backup? (yes/no): " CONFIRM
        if [[ "$CONFIRM" != "yes" ]]; then
            echo "Backup cancelled"
            return 1
        fi
    fi
    
    echo
    echo -e "${GREEN}Starting backup...${NC}"
    echo "=================="
    
    # Use dd for raw backup (like Clonezilla but simpler)
    if command -v pv >/dev/null 2>&1; then
        dd if="$SOURCE_DRIVE" bs=4M 2>/dev/null | pv -s "$SOURCE_SIZE" | dd of="$TARGET_DRIVE" bs=4M conv=fdatasync 2>/dev/null
    else
        dd if="$SOURCE_DRIVE" of="$TARGET_DRIVE" bs=4M status=progress conv=fdatasync
    fi
    
    RESULT=$?
    
    echo
    if [[ $RESULT -eq 0 ]]; then
        echo -e "${GREEN}"
        echo "================================================"
        echo "    BACKUP COMPLETED SUCCESSFULLY!"
        echo "================================================"
        echo -e "${NC}"
        echo "Your system has been backed up to $TARGET_DRIVE"
        echo "You can now use smart sync for future updates"
    else
        echo -e "${RED}"
        echo "================================================"
        echo "         BACKUP FAILED!"
        echo "================================================"
        echo -e "${NC}"
    fi
    
    echo
    echo "System will shutdown in 30 seconds..."
    echo "(Press any key to cancel shutdown)"
    
    if read -t 30 -n 1; then
        echo
        echo "Shutdown cancelled"
        echo "You can now use Clonezilla or run another backup"
    else
        echo
        echo "Shutting down..."
        sudo shutdown -h now
    fi
}

# Main execution
if [[ "$AUTO_MODE" == "true" ]]; then
    echo "🚀 AUTOMATIC BACKUP MODE"
    echo "========================="
    echo
fi

select_drives
perform_backup
EOF

sudo chmod +x "$DATA_MOUNT/auto-backup.sh"

# Create instructions
sudo tee "$DATA_MOUNT/README.txt" > /dev/null << 'EOF'
CLONEZILLA-BASED BACKUP USB
===========================

This USB combines Clonezilla Live with custom backup automation.

BOOT MENU OPTIONS:
🚀 Automatic System Backup - Boots directly to backup menu
🔧 Manual Backup Mode     - Access to both backup tools and Clonezilla
📦 Clonezilla Live        - Original Clonezilla functionality
🛠️  Clonezilla Expert     - Advanced Clonezilla options

AUTOMATIC MODE:
- Auto-detects drives
- Interactive drive selection
- Raw dd backup (like Clonezilla)
- Progress display
- Auto-shutdown when complete

MANUAL MODE:
- Access to shell
- Run: /media/user/BACKUP_TOOLS/auto-backup.sh
- Or use Clonezilla GUI

ADVANTAGES:
✅ Proven boot compatibility (Clonezilla)
✅ Professional disk cloning tools
✅ Both automatic and manual modes
✅ Raw disk backup for maximum compatibility
✅ Works on virtually any hardware

AFTER FIRST BACKUP:
Boot back to your normal system and use smart sync:
python3 backup_manager.py
EOF

# Cleanup
sudo umount "$ISO_MOUNT" "$BOOT_MOUNT" "$DATA_MOUNT"
sudo rmdir "$ISO_MOUNT" "$BOOT_MOUNT" "$DATA_MOUNT"

print_success "Clonezilla-based backup USB created!"
print_success "USB: $USB_DRIVE"
echo
print_success "FEATURES:"
print_success "✅ Clonezilla Live base (proven boot compatibility)"
print_success "✅ Custom backup automation"
print_success "✅ Automatic and manual modes"
print_success "✅ Professional disk cloning"
print_success "✅ Works on any hardware"
echo
print_warning "BOOT OPTIONS:"
print_warning "🚀 Automatic System Backup - Direct to backup"
print_warning "🔧 Manual Backup Mode - Shell + Clonezilla access"
print_warning "📦 Clonezilla Live - Original functionality"
echo
print_success "Ready to test! This should boot reliably in QEMU and real hardware."