#!/bin/bash
# System Backup Script - Command Line Version
# For use with cron jobs or manual execution

set -e

# Configuration
SOURCE_DRIVE=""              # Will be auto-detected
TARGET_DRIVE=""              # Will be detected or specified
RESTORE_MODE=false           # Restore mode flag
SYNC_MODE=false              # Smart sync mode flag
ANALYZE_ONLY=false           # Analysis only mode
LOG_FILE="/var/log/system_backup.log"

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

# Logging function
log() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    
    # Log to console
    echo "${timestamp} - ${message}"
    
    # Log to file if writable
    if [[ -w "$LOG_FILE" ]] || [[ -w "$(dirname "$LOG_FILE")" ]]; then
        echo "${timestamp} - ${message}" >> "$LOG_FILE" 2>/dev/null
    fi
}

# Error handling
error_exit() {
    log "${RED}ERROR: $1${NC}"
    exit 1
}

# Success message
success() {
    log "${GREEN}SUCCESS: $1${NC}"
}

# Warning message
warning() {
    log "${YELLOW}WARNING: $1${NC}"
}

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        error_exit "This script must be run as root (use sudo)"
    fi
}

# Detect root filesystem drive
detect_root_drive() {
    log "Detecting root filesystem drive..."
    
    # Find the device containing the root filesystem
    local root_device=$(df / | tail -1 | awk '{print $1}')
    
    # Remove partition number to get base device
    local base_device=$(echo "$root_device" | sed 's/[0-9]*$//')
    
    # Handle nvme drives (e.g., /dev/nvme0n1p1 -> /dev/nvme0n1)
    base_device=$(echo "$base_device" | sed 's/p$//')
    
    echo "$base_device"
}

# Check if target has existing backup
check_existing_backup() {
    local target_drive=$1
    local temp_mount="/tmp/backup_check_$$"
    
    log "Checking for existing backup on $target_drive..."
    
    # Get main partition
    local main_partition=$(lsblk -n -o NAME "$target_drive" | grep -v "^$(basename "$target_drive")$" | head -1)
    if [[ -z "$main_partition" ]]; then
        echo "false"
        return
    fi
    
    main_partition="/dev/$main_partition"
    
    # Try to mount and check
    mkdir -p "$temp_mount"
    if mount -o ro "$main_partition" "$temp_mount" 2>/dev/null; then
        if [[ -d "$temp_mount/etc" && -d "$temp_mount/home" && -d "$temp_mount/usr" ]]; then
            echo "true"
        else
            echo "false"
        fi
        umount "$temp_mount" 2>/dev/null
    else
        echo "false"
    fi
    rmdir "$temp_mount" 2>/dev/null
}

# Analyze changes between source and target
analyze_changes() {
    local source_drive=$1
    local target_drive=$2
    
    log "Analyzing changes between $source_drive and $target_drive..."
    
    # Check if target has existing backup
    local has_backup=$(check_existing_backup "$target_drive")
    
    if [[ "$has_backup" != "true" ]]; then
        log "No existing backup found. Full clone required."
        echo "FULL_CLONE_REQUIRED"
        return
    fi
    
    # Get filesystem usage for both drives
    local source_size=$(get_filesystem_size "$source_drive")
    local target_size=$(get_filesystem_size "$target_drive")
    
    # Calculate difference in GB
    local size_diff=$((${source_size} - ${target_size}))
    local size_diff_abs=${size_diff#-}  # Absolute value
    
    # Convert to GB (sizes are in KB)
    local size_diff_gb=$((size_diff_abs / 1024 / 1024))
    
    log "Source filesystem size: $((source_size / 1024 / 1024)) GB"
    log "Target filesystem size: $((target_size / 1024 / 1024)) GB"
    log "Size difference: ${size_diff_gb} GB"
    
    # Decision logic
    if [[ $size_diff_gb -lt 2 ]]; then
        log "Recommendation: Smart sync (minimal changes)"
        echo "SYNC_RECOMMENDED"
    elif [[ $size_diff_gb -lt 10 ]]; then
        log "Recommendation: Smart sync (moderate changes)"
        echo "SYNC_BENEFICIAL"
    else
        log "Recommendation: Full clone (major changes)"
        echo "FULL_CLONE_RECOMMENDED"
    fi
}

# Get filesystem size in KB
get_filesystem_size() {
    local drive=$1
    local temp_mount="/tmp/size_check_$$"
    
    # Get main partition
    local main_partition=$(lsblk -n -o NAME "$drive" | grep -v "^$(basename "$drive")$" | head -1)
    if [[ -z "$main_partition" ]]; then
        echo "0"
        return
    fi
    
    main_partition="/dev/$main_partition"
    
    # Mount and get usage
    mkdir -p "$temp_mount"
    if mount -o ro "$main_partition" "$temp_mount" 2>/dev/null; then
        local used_kb=$(df "$temp_mount" | tail -1 | awk '{print $3}')
        umount "$temp_mount" 2>/dev/null
        echo "$used_kb"
    else
        echo "0"
    fi
    rmdir "$temp_mount" 2>/dev/null
}

# Perform smart sync backup
smart_sync_backup() {
    local source=$1
    local target=$2
    
    log "Starting smart sync backup..."
    
    # Mount both filesystems
    local source_mount="/tmp/sync_source_$$"
    local target_mount="/tmp/sync_target_$$"
    
    mkdir -p "$source_mount" "$target_mount"
    
    # Get main partitions
    local source_partition=$(lsblk -n -o NAME "$source" | grep -v "^$(basename "$source")$" | head -1)
    local target_partition=$(lsblk -n -o NAME "$target" | grep -v "^$(basename "$target")$" | head -1)
    
    source_partition="/dev/$source_partition"
    target_partition="/dev/$target_partition"
    
    log "Mounting filesystems for sync..."
    mount -o ro "$source_partition" "$source_mount" || error_exit "Failed to mount source"
    mount "$target_partition" "$target_mount" || error_exit "Failed to mount target"
    
    # Perform rsync
    log "Starting rsync synchronization..."
    rsync -avHAXS \
        --numeric-ids \
        --delete \
        --progress \
        --exclude=/proc/* \
        --exclude=/sys/* \
        --exclude=/dev/* \
        --exclude=/tmp/* \
        --exclude=/run/* \
        --exclude=/mnt/* \
        --exclude=/media/* \
        --exclude=/lost+found \
        "$source_mount/" "$target_mount/" || {
        
        # Cleanup on failure
        umount "$source_mount" 2>/dev/null
        umount "$target_mount" 2>/dev/null
        rmdir "$source_mount" "$target_mount" 2>/dev/null
        error_exit "Smart sync failed"
    }
    
    # Cleanup
    umount "$source_mount" 2>/dev/null
    umount "$target_mount" 2>/dev/null
    rmdir "$source_mount" "$target_mount" 2>/dev/null
    
    success "Smart sync completed successfully!"
}

# Detect external drives
detect_external_drives() {
    log "Detecting external drives..."
    
    # Get all block devices
    lsblk -d -n -o NAME,SIZE,TYPE,TRAN | while read -r line; do
        if [[ $line == *"disk"* ]] && [[ $line == *"usb"* ]]; then
            drive_name=$(echo "$line" | awk '{print $1}')
            drive_size=$(echo "$line" | awk '{print $2}')
            echo "/dev/$drive_name ($drive_size)"
        fi
    done
}

# Validate drives
validate_drives() {
    if [[ ! -b "$SOURCE_DRIVE" ]]; then
        error_exit "Source drive $SOURCE_DRIVE does not exist or is not a block device"
    fi
    
    if [[ ! -b "$TARGET_DRIVE" ]]; then
        error_exit "Target drive $TARGET_DRIVE does not exist or is not a block device"
    fi
    
    if [[ "$SOURCE_DRIVE" == "$TARGET_DRIVE" ]]; then
        error_exit "Source and target drives cannot be the same"
    fi
    
    # Check if target drive is mounted
    if mount | grep -q "$TARGET_DRIVE"; then
        warning "Target drive $TARGET_DRIVE is currently mounted. Unmounting..."
        umount "$TARGET_DRIVE"* 2>/dev/null || true
    fi
}

# Get drive size
get_drive_size() {
    local drive=$1
    blockdev --getsize64 "$drive"
}

# Clone drive
clone_drive() {
    local source=$1
    local target=$2
    
    log "Starting drive clone operation..."
    log "Source: $source"
    log "Target: $target"
    
    # Get sizes
    source_size=$(get_drive_size "$source")
    target_size=$(get_drive_size "$target")
    
    log "Source size: $(numfmt --to=iec-i --suffix=B $source_size)"
    log "Target size: $(numfmt --to=iec-i --suffix=B $target_size)"
    
    if [[ $target_size -lt $source_size ]]; then
        error_exit "Target drive is smaller than source drive"
    fi
    
    # Start cloning
    log "Starting clone operation with dd..."
    
    if command -v pv >/dev/null 2>&1; then
        # Use pv for progress if available
        dd if="$source" bs=4M | pv -s "$source_size" | dd of="$target" bs=4M conv=fdatasync
    else
        # Use dd with status=progress
        dd if="$source" of="$target" bs=4M status=progress conv=fdatasync
    fi
    
    if [[ $? -eq 0 ]]; then
        success "Drive cloning completed successfully!"
        
        # Restore backup tools to external drive if this was a backup operation
        if [[ "$RESTORE_MODE" != true ]]; then
            log "Preserving backup tools on external drive..."
            local restore_script="$(dirname "$0")/restore_tools_after_backup.sh"
            if [[ -f "$restore_script" ]]; then
                "$restore_script" "$target" || log "Warning: Could not preserve backup tools"
            else
                log "Warning: Tool preservation script not found"
            fi
        fi
        
        # Sync to ensure all data is written
        log "Syncing data to disk..."
        sync
        
        # Verify partition table
        log "Verifying partition table on target drive..."
        fdisk -l "$target" | head -20 | tee -a "$LOG_FILE"
        
        success "Backup verification completed!"
    else
        error_exit "Drive cloning failed!"
    fi
}

# Create desktop entry
create_desktop_entry() {
    local desktop_file="$HOME/Desktop/System-Backup.desktop"
    local script_path=$(realpath "$0")
    
    cat > "$desktop_file" << EOF
[Desktop Entry]
Version=1.0
Type=Application
Name=System Backup
Comment=Clone internal drive to external M.2 SSD
Exec=gnome-terminal -- sudo "$script_path" --gui
Icon=drive-harddisk
Terminal=false
Categories=System;Utility;
EOF
    
    chmod +x "$desktop_file"
    log "Desktop entry created: $desktop_file"
}

# Show usage
show_usage() {
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  -s, --source DRIVE    Source drive (auto-detected if not specified)"
    echo "  -t, --target DRIVE    Target drive (required)"
    echo "  -r, --restore         Restore mode (reverse source and target)"
    echo "  --sync               Smart sync mode (faster incremental backup)"
    echo "  --analyze            Analyze changes without performing backup"
    echo "  -l, --list           List available drives"
    echo "  -d, --desktop        Create desktop entry"
    echo "  --gui               Launch GUI version"
    echo "  -h, --help          Show this help"
    echo ""
    echo "Examples:"
    echo "  $0 --list"
    echo "  $0 --analyze --target /dev/sdb"
    echo "  $0 --sync --target /dev/sdb"
    echo "  $0 --source /dev/nvme0n1 --target /dev/sdb"
    echo "  $0 --restore --source /dev/sdb --target /dev/nvme0n1"
    echo "  $0 --desktop"
    echo "  $0 --gui"
}

# Main function
main() {
    # Auto-detect source drive if not specified
    if [[ -z "$SOURCE_DRIVE" ]]; then
        SOURCE_DRIVE=$(detect_root_drive)
        log "Auto-detected root filesystem drive: $SOURCE_DRIVE"
    fi
    
    # Parse command line arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            -s|--source)
                SOURCE_DRIVE="$2"
                shift 2
                ;;
            -t|--target)
                TARGET_DRIVE="$2"
                shift 2
                ;;
            -r|--restore)
                RESTORE_MODE=true
                shift
                ;;
            --sync)
                SYNC_MODE=true
                shift
                ;;
            --analyze)
                ANALYZE_ONLY=true
                shift
                ;;
            -l|--list)
                echo "Available drives:"
                lsblk -d -o NAME,SIZE,TYPE,TRAN
                echo ""
                echo "External drives:"
                detect_external_drives
                exit 0
                ;;
            -d|--desktop)
                create_desktop_entry
                exit 0
                ;;
            --gui)
                python3 "$(dirname "$0")/backup_manager.py"
                exit 0
                ;;
            -h|--help)
                show_usage
                exit 0
                ;;
            *)
                error_exit "Unknown option: $1"
                ;;
        esac
    done
    
    # In restore mode, swap source and target for user convenience
    if [[ "$RESTORE_MODE" == true ]]; then
        if [[ -n "$SOURCE_DRIVE" && -n "$TARGET_DRIVE" ]]; then
            log "Restore mode: swapping source and target drives"
            TEMP_DRIVE="$SOURCE_DRIVE"
            SOURCE_DRIVE="$TARGET_DRIVE"
            TARGET_DRIVE="$TEMP_DRIVE"
        fi
    fi
    
    # Check if target drive is specified
    if [[ -z "$TARGET_DRIVE" ]]; then
        echo "Available external drives:"
        detect_external_drives
        echo ""
        read -p "Enter target drive (e.g., /dev/sdb): " TARGET_DRIVE
        
        if [[ -z "$TARGET_DRIVE" ]]; then
            error_exit "Target drive not specified"
        fi
    fi
    
    # Check root privileges
    check_root
    
    # Validate drives
    validate_drives
    
    # Handle analyze mode
    if [[ "$ANALYZE_ONLY" == "true" ]]; then
        echo ""
        echo "🔍 ANALYZING CHANGES"
        echo "Source: $SOURCE_DRIVE"
        echo "Target: $TARGET_DRIVE"
        echo ""
        
        analysis_result=$(analyze_changes "$SOURCE_DRIVE" "$TARGET_DRIVE")
        
        case "$analysis_result" in
            "FULL_CLONE_REQUIRED")
                echo "📋 ANALYSIS RESULT: Full Clone Required"
                echo "• No existing backup found on target drive"
                echo "• Complete drive cloning is necessary"
                ;;
            "SYNC_RECOMMENDED")
                echo "✅ ANALYSIS RESULT: Smart Sync Recommended"
                echo "• Minimal changes detected (< 2GB difference)"
                echo "• Smart sync will be much faster than full clone"
                ;;
            "SYNC_BENEFICIAL")
                echo "⚡ ANALYSIS RESULT: Smart Sync Beneficial"
                echo "• Moderate changes detected (< 10GB difference)"
                echo "• Smart sync recommended for faster backup"
                ;;
            "FULL_CLONE_RECOMMENDED")
                echo "🔄 ANALYSIS RESULT: Full Clone Recommended"
                echo "• Major changes detected (> 10GB difference)"
                echo "• Full clone may be more appropriate"
                ;;
        esac
        
        echo ""
        echo "Use --sync flag to perform smart sync backup"
        exit 0
    fi
    
    # Handle sync mode
    if [[ "$SYNC_MODE" == "true" ]]; then
        echo ""
        echo "⚡ SMART SYNC BACKUP"
        echo "Source: $SOURCE_DRIVE"
        echo "Target: $TARGET_DRIVE"
        echo ""
        
        # Check if sync is possible
        analysis_result=$(analyze_changes "$SOURCE_DRIVE" "$TARGET_DRIVE")
        
        if [[ "$analysis_result" == "FULL_CLONE_REQUIRED" ]]; then
            error_exit "Smart sync not possible: No existing backup found. Use full backup first."
        fi
        
        echo "Analysis: $analysis_result"
        echo ""
        read -p "Proceed with smart sync? (yes/no): " confirm
        if [[ "$confirm" != "yes" ]]; then
            log "Smart sync cancelled by user"
            exit 0
        fi
        
        smart_sync_backup "$SOURCE_DRIVE" "$TARGET_DRIVE"
        
        success "Smart sync backup completed successfully!"
        echo "Smart sync completed! External drive is up to date."
        exit 0
    fi
    
    # Confirm operation
    echo ""
    if [[ "$RESTORE_MODE" == true ]]; then
        echo "🚨 RESTORE CONFIGURATION 🚨"
        echo "This will RESTORE (overwrite target with source data):"
    else
        echo "BACKUP CONFIGURATION:"
        echo "This will BACKUP (copy source to target):"
    fi
    echo "Source: $SOURCE_DRIVE"
    echo "Target: $TARGET_DRIVE"
    echo ""
    echo "WARNING: All data on $TARGET_DRIVE will be DESTROYED!"
    
    if [[ "$RESTORE_MODE" == true ]]; then
        echo ""
        echo "⚠️  CRITICAL WARNING FOR RESTORE MODE ⚠️"
        echo "You are about to OVERWRITE $TARGET_DRIVE"
        echo "Make sure this is what you intend to do!"
        echo ""
        read -p "Type 'RESTORE' to confirm or anything else to cancel: " confirm
        if [[ "$confirm" != "RESTORE" ]]; then
            log "Restore operation cancelled by user"
            exit 0
        fi
    else
        read -p "Are you sure you want to continue? (yes/no): " confirm
        if [[ "$confirm" != "yes" ]]; then
            log "Operation cancelled by user"
            exit 0
        fi
    fi
    
    # Start operation
    if [[ "$RESTORE_MODE" == true ]]; then
        log "Starting system restore operation..."
    else
        log "Starting system backup operation..."
    fi
    clone_drive "$SOURCE_DRIVE" "$TARGET_DRIVE"
    
    log "System backup completed successfully!"
    echo ""
    echo "Backup completed! You can now safely remove the external drive."
}

# Run main function
main "$@"