#!/bin/bash
# Create Bootable Backup USB Script
# This creates a bootable USB that can perform system backups

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"
echo "==========================================="
echo

print_warning "This will create a bootable USB drive that can:"
print_warning "1. Boot into a minimal Linux environment"
print_warning "2. Automatically detect and backup your internal drive"
print_warning "3. Work completely offline (no running OS interference)"
echo

# List available USB drives
print_status "Available USB drives:"
lsblk -d -o NAME,SIZE,TYPE,TRAN | grep "usb" || {
    print_error "No USB drives detected!"
    exit 1
}

echo
read -p "Enter the USB drive to make bootable (e.g., /dev/sdb): " USB_DRIVE

if [[ ! -b "$USB_DRIVE" ]]; then
    print_error "Device $USB_DRIVE does not exist!"
    exit 1
fi

# Confirm USB drive selection
print_warning "⚠️  ALL DATA ON $USB_DRIVE WILL BE DESTROYED! ⚠️"
print_warning "This USB will become a bootable backup tool"
echo
lsblk "$USB_DRIVE"
echo
read -p "Are you sure you want to continue? (yes/no): " confirm

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

print_status "Creating bootable backup USB..."

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

# Create partition table and bootable partition
print_status "Creating partition table..."
sudo parted "$USB_DRIVE" --script mklabel gpt
sudo parted "$USB_DRIVE" --script mkpart ESP fat32 1MiB 512MiB
sudo parted "$USB_DRIVE" --script mkpart backup ext4 512MiB 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 "BOOT" "$BOOT_PART"
sudo mkfs.ext4 -L "BACKUP_TOOLS" "$DATA_PART"

# Mount partitions
BOOT_MOUNT="/tmp/usb_boot_$$"
DATA_MOUNT="/tmp/usb_data_$$"
sudo mkdir -p "$BOOT_MOUNT" "$DATA_MOUNT"
sudo mount "$BOOT_PART" "$BOOT_MOUNT"
sudo mount "$DATA_PART" "$DATA_MOUNT"

print_status "Installing backup tools..."

# Copy backup scripts to data partition
SCRIPT_DIR=$(dirname "$(realpath "$0")")
sudo cp -r "$SCRIPT_DIR"/* "$DATA_MOUNT/"
sudo chmod +x "$DATA_MOUNT"/*.sh

# Create autorun script for backup
sudo tee "$DATA_MOUNT/autorun_backup.sh" > /dev/null << 'EOF'
#!/bin/bash
# Auto-run backup script when booted from USB

set -e

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

echo -e "${GREEN}"
echo "========================================"
echo "   BOOTABLE BACKUP SYSTEM STARTED"
echo "========================================"
echo -e "${NC}"

# Wait for drives to be detected
sleep 5

echo "Detecting drives..."
echo

# Auto-detect internal drive (largest non-USB drive)
INTERNAL_DRIVE=$(lsblk -d -n -o NAME,SIZE,TYPE,TRAN | grep "disk" | grep -v "usb" | sort -k2 -hr | head -1 | awk '{print "/dev/" $1}')

# Auto-detect backup target (largest USB drive that's not the boot drive)
BOOT_USB=$(df /backup-tools | tail -1 | awk '{print $1}' | sed 's/[0-9]*$//')
TARGET_DRIVE=$(lsblk -d -n -o NAME,SIZE,TYPE,TRAN | grep "disk" | grep "usb" | awk '{print "/dev/" $1}' | grep -v "$BOOT_USB" | head -1)

echo "Auto-detected configuration:"
echo "Internal drive: $INTERNAL_DRIVE"
echo "Target drive: $TARGET_DRIVE"
echo

if [[ -z "$INTERNAL_DRIVE" || -z "$TARGET_DRIVE" ]]; then
    echo -e "${RED}Could not auto-detect drives. Manual selection required.${NC}"
    echo "Available drives:"
    lsblk -d -o NAME,SIZE,TYPE,TRAN
    echo
    read -p "Enter source drive: " INTERNAL_DRIVE
    read -p "Enter target drive: " TARGET_DRIVE
fi

echo -e "${YELLOW}FINAL CONFIRMATION${NC}"
echo "Source (internal): $INTERNAL_DRIVE"
echo "Target (backup): $TARGET_DRIVE"
echo
echo -e "${RED}⚠️  TARGET DRIVE WILL BE COMPLETELY OVERWRITTEN! ⚠️${NC}"
echo
read -p "Continue with backup? (yes/no): " confirm

if [[ "$confirm" != "yes" ]]; then
    echo "Backup cancelled"
    exit 0
fi

echo -e "${GREEN}Starting backup...${NC}"
cd /backup-tools
./backup_script.sh --source "$INTERNAL_DRIVE" --target "$TARGET_DRIVE"

echo -e "${GREEN}"
echo "========================================"
echo "   BACKUP COMPLETED SUCCESSFULLY!"
echo "========================================"
echo -e "${NC}"
echo "You can now:"
echo "1. Remove the USB drives"
echo "2. Reboot to your normal system"
echo "3. Use smart sync for future backups"
echo
read -p "Press Enter to shutdown system..."
shutdown -h now
EOF

sudo chmod +x "$DATA_MOUNT/autorun_backup.sh"

print_status "Downloading minimal Linux for USB boot..."

# Check if we have a Linux ISO or create a simple boot setup
print_warning "Note: You'll need to manually add a bootable Linux distro"
print_warning "Recommendation: Use Ubuntu Live USB creator or similar"
print_warning "Then copy the backup tools to the USB"

# Create instructions file
sudo tee "$DATA_MOUNT/INSTRUCTIONS.txt" > /dev/null << EOF
BOOTABLE BACKUP USB INSTRUCTIONS
==================================

This USB contains backup tools but needs a bootable Linux environment.

SETUP STEPS:
1. Use Ubuntu's "Startup Disk Creator" or similar tool
2. Create a bootable Ubuntu Live USB on this drive
3. Boot from this USB
4. Open terminal and run:
   cd /media/ubuntu/BACKUP_TOOLS
   sudo ./autorun_backup.sh

AUTOMATIC BACKUP:
- The script will auto-detect your internal drive
- Auto-detect external backup target
- Perform complete system backup
- Shutdown when complete

AFTER FIRST BACKUP:
- Boot back to your normal system
- Use the GUI for smart sync backups:
  python3 backup_manager.py
- Click "Smart Sync Backup" for fast updates

MANUAL BACKUP:
If auto-detection fails, you can run manually:
sudo ./backup_script.sh --source /dev/nvme0n1 --target /dev/sda
EOF

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

print_success "Bootable backup USB preparation complete!"
print_success "USB: $USB_DRIVE"
echo
print_warning "NEXT STEPS:"
print_warning "1. Use Ubuntu's 'Startup Disk Creator' to make this USB bootable"
print_warning "2. Boot from USB in BIOS/UEFI boot menu"
print_warning "3. Run: sudo /media/ubuntu/BACKUP_TOOLS/autorun_backup.sh"
echo
print_success "After first backup, use GUI smart sync for incremental updates!"