#!/bin/bash
# Portable Backup Tool Installer for External M.2 SSD
# This script sets up the backup tools on the external drive so they survive cloning operations

set -e

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

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

# Detect external drives
detect_external_drive() {
    print_status "Detecting external M.2 SSD..."
    
    # Look for USB-connected drives
    local external_drives=()
    while IFS= read -r line; do
        if [[ $line == *"disk"* ]] && [[ $line == *"usb"* ]]; then
            local drive_name=$(echo "$line" | awk '{print $1}')
            local drive_size=$(echo "$line" | awk '{print $4}')
            external_drives+=("/dev/$drive_name ($drive_size)")
        fi
    done < <(lsblk -d -o NAME,SIZE,TYPE,TRAN | grep -E "disk.*usb")
    
    if [[ ${#external_drives[@]} -eq 0 ]]; then
        print_error "No external USB drives found. Please connect your M.2 SSD."
        exit 1
    fi
    
    if [[ ${#external_drives[@]} -eq 1 ]]; then
        EXTERNAL_DRIVE=$(echo "${external_drives[0]}" | cut -d' ' -f1)
        print_success "Auto-detected external drive: ${external_drives[0]}"
    else
        print_status "Multiple external drives found:"
        for i in "${!external_drives[@]}"; do
            echo "  $((i+1)). ${external_drives[i]}"
        done
        read -p "Select drive number: " selection
        EXTERNAL_DRIVE=$(echo "${external_drives[$((selection-1))]}" | cut -d' ' -f1)
    fi
}

# Create backup tools partition
create_tools_partition() {
    print_status "Setting up backup tools partition on $EXTERNAL_DRIVE..."
    
    # Check if there's already a small partition at the end
    local last_partition=$(lsblk -n -o NAME "$EXTERNAL_DRIVE" | tail -1)
    local tools_partition="${EXTERNAL_DRIVE}p3"  # Assuming nvme, adjust for sda
    
    # If it's sda style, adjust
    if [[ $EXTERNAL_DRIVE == *"sda"* ]]; then
        tools_partition="${EXTERNAL_DRIVE}3"
    fi
    
    # Check if tools partition already exists
    if lsblk | grep -q "$(basename "$tools_partition")"; then
        print_warning "Tools partition already exists: $tools_partition"
        TOOLS_PARTITION="$tools_partition"
        return
    fi
    
    print_warning "This will create a 512MB partition at the end of $EXTERNAL_DRIVE"
    print_warning "This will slightly reduce the cloneable space but preserve backup tools"
    read -p "Continue? (yes/no): " confirm
    
    if [[ "$confirm" != "yes" ]]; then
        print_error "Operation cancelled"
        exit 1
    fi
    
    # Create 512MB partition at the end using parted
    sudo parted "$EXTERNAL_DRIVE" --script mkpart primary ext4 -512MiB 100%
    
    # Get the new partition name
    TOOLS_PARTITION=$(lsblk -n -o NAME "$EXTERNAL_DRIVE" | tail -1)
    TOOLS_PARTITION="/dev/$TOOLS_PARTITION"
    
    # Format the partition
    sudo mkfs.ext4 -L "BACKUP_TOOLS" "$TOOLS_PARTITION"
    
    print_success "Tools partition created: $TOOLS_PARTITION"
}

# Install backup tools to external drive
install_tools_to_external() {
    local mount_point="/mnt/backup_tools"
    
    print_status "Installing backup tools to external drive..."
    
    # Create mount point
    sudo mkdir -p "$mount_point"
    
    # Mount tools partition
    sudo mount "$TOOLS_PARTITION" "$mount_point"
    
    # Copy all backup tools
    sudo cp -r . "$mount_point/backup_system"
    
    # Create launcher script that works from any location
    sudo tee "$mount_point/backup_system/launch_backup_tools.sh" > /dev/null << 'EOF'
#!/bin/bash
# Portable launcher for backup tools

# Get the directory where this script is located
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Make sure scripts are executable
chmod +x *.sh *.py

# Launch GUI if X11 is available, otherwise show CLI options
if [[ -n "$DISPLAY" ]]; then
    echo "Launching Backup Manager GUI..."
    python3 backup_manager.py
else
    echo "No GUI available. Command line options:"
    ./backup_script.sh --help
fi
EOF
    
    sudo chmod +x "$mount_point/backup_system/launch_backup_tools.sh"
    
    # Create desktop entry for when booted from external drive
    sudo tee "$mount_point/backup_system/create_desktop_entry.sh" > /dev/null << 'EOF'
#!/bin/bash
# Create desktop entry when booted from external drive

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DESKTOP_DIR="$HOME/Desktop"
APPLICATIONS_DIR="$HOME/.local/share/applications"

mkdir -p "$APPLICATIONS_DIR"

# Create desktop entry
cat > "$APPLICATIONS_DIR/portable-backup.desktop" << EOL
[Desktop Entry]
Version=1.0
Type=Application
Name=Portable Backup Manager
Comment=Restore internal drive from external backup
Exec=python3 "$SCRIPT_DIR/backup_manager.py"
Icon=drive-harddisk
Terminal=false
Categories=System;Utility;
StartupNotify=true
EOL

# Create desktop shortcut if Desktop exists
if [[ -d "$DESKTOP_DIR" ]]; then
    cp "$APPLICATIONS_DIR/portable-backup.desktop" "$DESKTOP_DIR/"
    chmod +x "$DESKTOP_DIR/portable-backup.desktop"
fi

echo "Desktop entry created for portable backup tools"
EOF
    
    sudo chmod +x "$mount_point/backup_system/create_desktop_entry.sh"
    
    # Unmount
    sudo umount "$mount_point"
    
    print_success "Backup tools installed to external drive"
}

# Create post-backup restoration script
create_restoration_script() {
    print_status "Creating post-backup tool restoration script..."
    
    # This script will be called after each backup to restore the tools
    cat > "restore_tools_after_backup.sh" << 'EOF'
#!/bin/bash
# Restore backup tools after cloning operation
# This script runs automatically after backup to preserve tools on external drive

set -e

print_status() {
    echo "[$(date '+%H:%M:%S')] $1"
}

# Detect the external drive (should be the target of backup)
EXTERNAL_DRIVE="$1"
if [[ -z "$EXTERNAL_DRIVE" ]]; then
    print_status "Auto-detecting external drive..."
    EXTERNAL_DRIVE=$(lsblk -d -o NAME,TRAN | grep usb | awk '{print "/dev/" $1}' | head -1)
fi

if [[ -z "$EXTERNAL_DRIVE" ]]; then
    echo "Error: Could not detect external drive"
    exit 1
fi

print_status "Restoring backup tools to $EXTERNAL_DRIVE"

# Find the tools partition (should be the last partition)
TOOLS_PARTITION=$(lsblk -n -o NAME "$EXTERNAL_DRIVE" | tail -1)
TOOLS_PARTITION="/dev/$TOOLS_PARTITION"

# Check if tools partition exists
if ! lsblk | grep -q "$(basename "$TOOLS_PARTITION")"; then
    echo "Warning: Tools partition not found. Backup tools may have been overwritten."
    exit 1
fi

# Mount tools partition
MOUNT_POINT="/mnt/backup_tools_restore"
mkdir -p "$MOUNT_POINT"
mount "$TOOLS_PARTITION" "$MOUNT_POINT"

# Check if tools exist
if [[ -d "$MOUNT_POINT/backup_system" ]]; then
    print_status "Backup tools preserved on external drive"
    
    # Update the tools with any changes from source
    rsync -av --exclude='.git' "$(dirname "$0")/" "$MOUNT_POINT/backup_system/"
    
    print_status "Backup tools updated on external drive"
else
    print_status "Installing backup tools to external drive for first time"
    cp -r "$(dirname "$0")" "$MOUNT_POINT/backup_system"
fi

# Ensure scripts are executable
chmod +x "$MOUNT_POINT/backup_system"/*.sh
chmod +x "$MOUNT_POINT/backup_system"/*.py

umount "$MOUNT_POINT"
print_status "Backup tools restoration complete"
EOF
    
    chmod +x "restore_tools_after_backup.sh"
    
    print_success "Post-backup restoration script created"
}

# Update backup scripts to call restoration
update_backup_scripts() {
    print_status "Updating backup scripts to preserve tools..."
    
    # Add tool restoration to GUI backup manager
    if ! grep -q "restore_tools_after_backup" backup_manager.py; then
        # Add restoration call after successful backup
        sed -i '/self\.log("Backup completed successfully!")/a\\n            # Restore backup tools to external drive\n            try:\n                subprocess.run([os.path.join(os.path.dirname(__file__), "restore_tools_after_backup.sh"), target], check=False)\n            except Exception as e:\n                self.log(f"Warning: Could not restore tools to external drive: {e}")' backup_manager.py
    fi
    
    # Add tool restoration to command line script
    if ! grep -q "restore_tools_after_backup" backup_script.sh; then
        sed -i '/success "Drive cloning completed successfully!"/a\\n    # Restore backup tools to external drive\n    log "Restoring backup tools to external drive..."\n    if [[ -f "$(dirname "$0")/restore_tools_after_backup.sh" ]]; then\n        "$(dirname "$0")/restore_tools_after_backup.sh" "$target" || log "Warning: Could not restore tools"\n    fi' backup_script.sh
    fi
    
    print_success "Backup scripts updated to preserve tools"
}

# Create auto-mount script for tools partition
create_automount_script() {
    print_status "Creating auto-mount script for backup tools..."
    
    cat > "mount_backup_tools.sh" << 'EOF'
#!/bin/bash
# Auto-mount backup tools partition and create desktop shortcut

# Find tools partition by label
TOOLS_PARTITION=$(blkid -L "BACKUP_TOOLS" 2>/dev/null)

if [[ -z "$TOOLS_PARTITION" ]]; then
    echo "Backup tools partition not found"
    exit 1
fi

# Create mount point
MOUNT_POINT="$HOME/backup_tools"
mkdir -p "$MOUNT_POINT"

# Mount if not already mounted
if ! mountpoint -q "$MOUNT_POINT"; then
    mount "$TOOLS_PARTITION" "$MOUNT_POINT" 2>/dev/null || {
        echo "Mounting with sudo..."
        sudo mount "$TOOLS_PARTITION" "$MOUNT_POINT"
    }
fi

echo "Backup tools mounted at: $MOUNT_POINT"

# Create desktop entry if tools exist
if [[ -d "$MOUNT_POINT/backup_system" ]]; then
    cd "$MOUNT_POINT/backup_system"
    ./create_desktop_entry.sh
    echo "Desktop entry created for backup tools"
fi
EOF
    
    chmod +x "mount_backup_tools.sh"
    
    print_success "Auto-mount script created"
}

# Main installation
main() {
    echo ""
    echo "=============================================="
    echo "  Portable Backup Tools Installer"
    echo "  For External M.2 SSD"
    echo "=============================================="
    echo ""
    
    print_warning "This installer will:"
    print_warning "1. Create a 512MB tools partition on your external M.2 SSD"
    print_warning "2. Install backup tools that survive cloning operations"
    print_warning "3. Set up automatic tool restoration after backups"
    print_warning "4. Enable booting from external drive with restore capability"
    echo ""
    
    read -p "Continue? (yes/no): " confirm
    if [[ "$confirm" != "yes" ]]; then
        print_error "Installation cancelled"
        exit 1
    fi
    
    detect_external_drive
    create_tools_partition
    install_tools_to_external
    create_restoration_script
    update_backup_scripts
    create_automount_script
    
    echo ""
    print_success "Portable backup tools installation complete!"
    echo ""
    echo "Your external M.2 SSD now has:"
    echo "  • Preserved backup tools in separate partition"
    echo "  • Automatic tool restoration after each backup"
    echo "  • Bootable system restoration capability"
    echo ""
    echo "When booted from external drive:"
    echo "  • Run: ~/backup_tools/backup_system/launch_backup_tools.sh"
    echo "  • Or use desktop shortcut if available"
    echo ""
    print_warning "Note: Your external drive now has 512MB less space for cloning"
    print_warning "But the backup tools will always be available for system recovery!"
}

# 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

main "$@"