#!/bin/bash
# Restore backup tools after cloning operation
# This script runs automatically after backup to preserve tools on external drive

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}[$(date '+%H:%M:%S')]${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"
}

# Get the external drive from parameter or auto-detect
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
    print_error "Could not detect external drive"
    exit 1
fi

print_status "Restoring backup tools to $EXTERNAL_DRIVE"

# Find the tools partition (look for BACKUP_TOOLS label first)
TOOLS_PARTITION=$(blkid -L "BACKUP_TOOLS" 2>/dev/null || echo "")

# If not found by label, try to find the last partition on the external drive
if [[ -z "$TOOLS_PARTITION" ]]; then
    # Get all partitions on the external drive
    mapfile -t partitions < <(lsblk -n -o NAME "$EXTERNAL_DRIVE" | grep -v "^$(basename "$EXTERNAL_DRIVE")$")
    
    if [[ ${#partitions[@]} -gt 0 ]]; then
        # Check the last partition
        last_partition="/dev/${partitions[-1]}"
        
        # Check if it's small (likely our tools partition)
        partition_size=$(lsblk -n -o SIZE "$last_partition" | tr -d ' ')
        if [[ "$partition_size" == *M ]] && [[ ${partition_size%M} -le 1024 ]]; then
            TOOLS_PARTITION="$last_partition"
            print_status "Found potential tools partition: $TOOLS_PARTITION ($partition_size)"
        fi
    fi
fi

# If still no tools partition found, create one
if [[ -z "$TOOLS_PARTITION" ]]; then
    print_warning "No tools partition found. Creating one..."
    
    # Create 512MB partition at the end
    if command -v parted >/dev/null 2>&1; then
        parted "$EXTERNAL_DRIVE" --script mkpart primary ext4 -512MiB 100% || {
            print_warning "Could not create partition. Backup tools will not be preserved."
            exit 0
        }
        
        # Wait for partition to appear
        sleep 2
        
        # Get the new partition
        mapfile -t partitions < <(lsblk -n -o NAME "$EXTERNAL_DRIVE" | grep -v "^$(basename "$EXTERNAL_DRIVE")$")
        TOOLS_PARTITION="/dev/${partitions[-1]}"
        
        # Format it
        mkfs.ext4 -L "BACKUP_TOOLS" "$TOOLS_PARTITION" -F >/dev/null 2>&1 || {
            print_warning "Could not format tools partition"
            exit 0
        }
        
        print_success "Created tools partition: $TOOLS_PARTITION"
    else
        print_warning "parted not available. Cannot create tools partition."
        exit 0
    fi
fi

# Mount tools partition
MOUNT_POINT="/tmp/backup_tools_restore_$$"
mkdir -p "$MOUNT_POINT"

mount "$TOOLS_PARTITION" "$MOUNT_POINT" 2>/dev/null || {
    print_error "Could not mount tools partition $TOOLS_PARTITION"
    rmdir "$MOUNT_POINT"
    exit 1
}

# Ensure we unmount on exit
trap 'umount "$MOUNT_POINT" 2>/dev/null; rmdir "$MOUNT_POINT" 2>/dev/null' EXIT

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

# Install/update backup tools
if [[ -d "$MOUNT_POINT/backup_system" ]]; then
    print_status "Updating existing backup tools..."
    # Sync changes, excluding git and cache files
    rsync -av --delete --exclude='.git' --exclude='__pycache__' --exclude='*.pyc' \
          "$SCRIPT_DIR/" "$MOUNT_POINT/backup_system/"
else
    print_status "Installing backup tools for first time..."
    mkdir -p "$MOUNT_POINT/backup_system"
    cp -r "$SCRIPT_DIR"/* "$MOUNT_POINT/backup_system/"
fi

# Create portable launcher if it doesn't exist
if [[ ! -f "$MOUNT_POINT/backup_system/launch_backup_tools.sh" ]]; then
    cat > "$MOUNT_POINT/backup_system/launch_backup_tools.sh" << '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

echo "=========================================="
echo "  Portable Backup Tools"
echo "=========================================="
echo ""
echo "Available options:"
echo "1. Launch GUI Backup Manager"
echo "2. Command Line Backup"
echo "3. Command Line Restore"
echo "4. List Available Drives"
echo "5. Create Desktop Entry"
echo ""

# Check if GUI is available
if [[ -n "$DISPLAY" ]] && command -v python3 >/dev/null 2>&1; then
    read -p "Select option (1-5): " choice
    case $choice in
        1)
            echo "Launching GUI Backup Manager..."
            python3 backup_manager.py
            ;;
        2)
            echo "Starting command line backup..."
            ./backup_script.sh
            ;;
        3)
            echo "Starting command line restore..."
            ./backup_script.sh --restore
            ;;
        4)
            echo "Available drives:"
            ./backup_script.sh --list
            ;;
        5)
            ./create_desktop_entry.sh
            ;;
        *)
            echo "Invalid option"
            ;;
    esac
else
    echo "GUI not available. Use command line options:"
    ./backup_script.sh --help
fi
EOF
fi

# Create desktop entry creator
cat > "$MOUNT_POINT/backup_system/create_desktop_entry.sh" << 'EOF'
#!/bin/bash
# Create desktop entry for portable backup tools

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

mkdir -p "$APPLICATIONS_DIR"

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

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

echo "Application entry created: $APPLICATIONS_DIR/portable-backup.desktop"
EOF

# Make all scripts executable
chmod +x "$MOUNT_POINT/backup_system"/*.sh
chmod +x "$MOUNT_POINT/backup_system"/*.py

# Create a README for the external drive
cat > "$MOUNT_POINT/backup_system/README_EXTERNAL.md" << 'EOF'
# Portable Backup Tools

This external drive contains both:
1. **Your system backup** (main partitions)
2. **Backup tools** (this partition)

## When Booted From This External Drive:

### Quick Start:
```bash
# Mount tools and launch
sudo mkdir -p /mnt/tools
sudo mount LABEL=BACKUP_TOOLS /mnt/tools
cd /mnt/tools/backup_system
./launch_backup_tools.sh
```

### Or Create Desktop Entry:
```bash
cd /mnt/tools/backup_system
./create_desktop_entry.sh
```

## Common Operations:

### Restore Internal Drive:
1. Boot from this external drive
2. Launch backup tools
3. Select "Restore from External"
4. Choose external → internal
5. Click "Reboot & Restore"

### Update Backup:
1. Boot normally from internal drive
2. Connect this external drive
3. Run backup as usual
4. Tools will be automatically preserved

## Drive Layout:
- Partition 1-2: System backup (bootable)
- Last Partition: Backup tools (this)
EOF

print_success "Backup tools preserved on external drive"
print_status "Tools available at: $TOOLS_PARTITION"
print_status "To access when booted from external: mount LABEL=BACKUP_TOOLS /mnt/tools"

exit 0