#!/bin/bash

# Live System Preparation Script for LVM Migration
# This script prepares a live USB system with all necessary tools for LVM migration

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

log() {
    echo -e "${BLUE}[$(date '+%H:%M:%S')]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1" >&2
    exit 1
}

warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

check_live_system() {
    log "Checking if running from live system..."
    
    # Check multiple indicators of live system
    local is_live=false
    
    # Check for common live system indicators
    if [ -f "/etc/casper.conf" ] || [ -f "/lib/live/mount/medium" ]; then
        is_live=true
    fi
    
    # Check if root filesystem is on loop, overlay, or tmpfs
    local root_device=$(df / | tail -1 | awk '{print $1}')
    if [[ "$root_device" == *"loop"* ]] || [[ "$root_device" == *"overlay"* ]] || [[ "$root_device" == *"tmpfs"* ]]; then
        is_live=true
    fi
    
    # Check for live system processes
    if pgrep -f "casper" >/dev/null 2>&1 || pgrep -f "live-boot" >/dev/null 2>&1; then
        is_live=true
    fi
    
    if [ "$is_live" = true ]; then
        success "Confirmed: Running from live system"
    else
        warning "This doesn't appear to be a live system!"
        echo "For safety, LVM migration should be run from a live USB system."
        read -p "Continue anyway? [y/N] " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            error "Operation aborted. Please boot from a live USB system."
        fi
    fi
}

update_package_lists() {
    log "Updating package lists..."
    
    if command -v apt >/dev/null 2>&1; then
        apt update || warning "Failed to update apt package lists"
    elif command -v pacman >/dev/null 2>&1; then
        pacman -Sy || warning "Failed to update pacman package lists"
    elif command -v dnf >/dev/null 2>&1; then
        dnf makecache || warning "Failed to update dnf package cache"
    else
        warning "Unknown package manager - manual tool installation may be required"
    fi
}

install_required_tools() {
    log "Installing required tools for LVM migration..."
    
    # Detect distribution for package name variations
    local distro="unknown"
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        distro="$ID"
        log "Detected distribution: $PRETTY_NAME"
    fi
    
    # Define package groups with alternatives for different distributions
    local package_groups=(
        "lvm:lvm2,lvm"
        "cryptsetup:cryptsetup,cryptsetup-bin"
        "rsync:rsync"
        "parted:parted"
        "pv:pv,pipe-viewer"
        "grub-efi:grub-efi-amd64,grub-efi,grub-efi-amd64-bin"
        "grub-pc:grub-pc-bin,grub-pc"
        "grub-common:grub-common,grub2-common"
        "e2fsprogs:e2fsprogs"
        "dosfstools:dosfstools,mtools"
        "util-linux:util-linux"
        "coreutils:coreutils"
        "bc:bc"
        "initramfs:initramfs-tools,dracut"
        "udev:udev,systemd-udev"
        "kmod:kmod,module-init-tools"
    )
    
    local missing_tools=()
    local packages_to_install=()
    
    # Check which tools are missing
    for tool_spec in "${tools_to_check[@]}"; do
        local tool=$(echo "$tool_spec" | cut -d: -f1)
        local packages=$(echo "$tool_spec" | cut -d: -f2)
        
        if ! command -v "$tool" >/dev/null 2>&1; then
            missing_tools+=("$tool")
            # Add packages (handle comma-separated list)
            IFS=',' read -ra PKGS <<< "$packages"
            for pkg in "${PKGS[@]}"; do
                if [[ ! " ${packages_to_install[@]} " =~ " ${pkg} " ]]; then
                    packages_to_install+=("$pkg")
                fi
            done
        fi
    done
    
    if [ ${#missing_tools[@]} -eq 0 ]; then
        success "All required tools are already available"
        return
    fi
    
    echo "Missing tools: ${missing_tools[*]}"
    echo "Will attempt to install: ${packages_to_install[*]}"
    
    # Install packages based on available package manager
    if command -v apt >/dev/null 2>&1; then
        log "Installing packages with apt..."
        
        # Update package lists first
        apt update || warning "Failed to update package lists"
        
        # Function to try installing packages with alternatives
        try_install() {
            local desc="$1"
            local packages_str="$2"
            IFS=',' read -ra packages <<< "$packages_str"
            
            log "Installing $desc..."
            for pkg in "${packages[@]}"; do
                if apt install -y "$pkg" >/dev/null 2>&1; then
                    success "Installed $pkg for $desc"
                    return 0
                fi
            done
            warning "Failed to install any package for $desc (tried: ${packages_str//,/, })"
            return 1
        }
        
        # Install packages by groups
        for group in "${package_groups[@]}"; do
            local desc="${group%:*}"
            local packages="${group#*:}"
            try_install "$desc" "$packages"
        done
    elif command -v pacman >/dev/null 2>&1; then
        log "Installing packages with pacman..."
        pacman -S --noconfirm "${packages_to_install[@]}" || {
            warning "Some packages failed to install via pacman"
        }
    elif command -v dnf >/dev/null 2>&1; then
        log "Installing packages with dnf..."
        dnf install -y "${packages_to_install[@]}" || {
            warning "Some packages failed to install via dnf"
        }
    else
        warning "Unknown package manager. Please install these packages manually:"
        echo "  ${packages_to_install[*]}"
    fi
    
    # Verify installation
    local still_missing=()
    for tool_spec in "${tools_to_check[@]}"; do
        local tool=$(echo "$tool_spec" | cut -d: -f1)
        if ! command -v "$tool" >/dev/null 2>&1; then
            still_missing+=("$tool")
        fi
    done
    
    if [ ${#still_missing[@]} -eq 0 ]; then
        success "All required tools are now available"
    else
        error "Still missing tools: ${still_missing[*]}. Please install them manually."
    fi
}

enable_lvm_kernel_modules() {
    log "Enabling LVM kernel modules..."
    
    local modules=("dm_mod" "dm_crypt" "dm_snapshot")
    
    for module in "${modules[@]}"; do
        if ! lsmod | grep -q "^$module"; then
            log "Loading kernel module: $module"
            modprobe "$module" || warning "Failed to load $module module"
        else
            log "Module $module already loaded"
        fi
    done
    
    # Start LVM services if available
    if command -v systemctl >/dev/null 2>&1; then
        systemctl start lvm2-monitor 2>/dev/null || true
        systemctl start lvm2-lvmetad 2>/dev/null || true
    fi
    
    success "LVM kernel modules enabled"
}

check_drive_availability() {
    log "Checking for available drives..."
    
    # List all available block devices
    echo "Available drives:"
    lsblk -dpno NAME,SIZE,MODEL,VENDOR | grep -E "sd[a-z]|nvme[0-9]|mmcblk[0-9]" | while read line; do
        echo "  $line"
    done
    
    # Count drives
    local drive_count=$(lsblk -dpno NAME | grep -E "sd[a-z]|nvme[0-9]|mmcblk[0-9]" | wc -l)
    
    if [ "$drive_count" -lt 2 ]; then
        warning "Only $drive_count drive(s) detected. Migration requires at least 2 drives:"
        echo "  1. Internal drive (source)"
        echo "  2. External M.2 SSD (target)"
        echo
        echo "Please connect your external M.2 SSD and try again."
        exit 1
    else
        success "Found $drive_count drives - sufficient for migration"
    fi
}

create_migration_workspace() {
    log "Creating migration workspace..."
    
    local workspace="/tmp/lvm-migration"
    mkdir -p "$workspace"/{scripts,logs,mounts}
    
    # Copy migration script to workspace if it exists
    if [ -f "./migrate_to_lvm.sh" ]; then
        cp "./migrate_to_lvm.sh" "$workspace/scripts/"
        chmod +x "$workspace/scripts/migrate_to_lvm.sh"
        success "Migration script copied to workspace"
    fi
    
    # Create useful aliases
    cat > "$workspace/aliases.sh" << 'EOF'
#!/bin/bash
# Useful aliases for LVM migration

alias ll='ls -la'
alias drives='lsblk -dpno NAME,SIZE,MODEL,VENDOR'
alias mounts='mount | grep -E "sd[a-z]|nvme|loop|mapper"'
alias lvminfo='pvs && vgs && lvs'
alias migration='cd /tmp/lvm-migration && ls -la'

echo "Migration workspace aliases loaded:"
echo "  drives  - List all available drives"
echo "  mounts  - Show mounted filesystems"
echo "  lvminfo - Display LVM information"
echo "  migration - Go to migration workspace"
EOF
    
    success "Migration workspace created at $workspace"
    echo "To load helpful aliases: source $workspace/aliases.sh"
}

main() {
    echo -e "${GREEN}=== Live System Preparation for LVM Migration ===${NC}"
    echo "This script prepares your live USB system for LVM migration"
    echo
    
    # Check if running as root
    if [ "$EUID" -ne 0 ]; then
        error "This script must be run as root. Use: sudo $0"
    fi
    
    check_live_system
    update_package_lists
    install_required_tools
    enable_lvm_kernel_modules
    check_drive_availability
    create_migration_workspace
    
    success "Live system preparation completed!"
    echo
    echo -e "${GREEN}Next steps:${NC}"
    echo "1. Connect your external M.2 SSD if not already connected"
    echo "2. Run the LVM migration script:"
    echo "   sudo ./migrate_to_lvm.sh"
    echo "3. Follow the interactive prompts to complete migration"
    echo
    echo -e "${YELLOW}Important reminders:${NC}"
    echo "• Ensure your external M.2 SSD is large enough for your system"
    echo "• The migration will DESTROY all data on the target drive"
    echo "• Your original internal drive will remain unchanged as backup"
    echo "• After migration, update BIOS boot order to boot from external drive"
}

main "$@"