#!/bin/bash

# Enhanced LVM Snapshot Backup System
# Integrates LVM snapshots with the existing backup system for M.2 external drives

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

# Configuration
VG_NAME="system-vg"
SNAPSHOT_SIZE="10G"
BACKUP_BASE_DIR="/mnt/backup"
EXTERNAL_BACKUP_DRIVE=""  # Will be auto-detected
LOG_FILE="/var/log/lvm-snapshot-backup.log"

# Snapshot names
ROOT_SNAPSHOT="root-snapshot"
HOME_SNAPSHOT="home-snapshot"
BOOT_SNAPSHOT="boot-snapshot"

log() {
    local message="[$(date '+%Y-%m-%d %H:%M:%S')] $1"
    echo -e "${BLUE}$message${NC}"
    echo "$message" >> "$LOG_FILE"
}

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

warning() {
    local message="[WARNING] $1"
    echo -e "${YELLOW}$message${NC}"
    echo "$message" >> "$LOG_FILE"
}

success() {
    local message="[SUCCESS] $1"
    echo -e "${GREEN}$message${NC}"
    echo "$message" >> "$LOG_FILE"
}

check_lvm_system() {
    log "Checking LVM system status..."
    
    # Check if we're running on an LVM system
    if ! command -v lvm >/dev/null 2>&1; then
        error "LVM tools not found. This system may not support LVM."
    fi
    
    # Check if volume group exists
    if ! vgs "$VG_NAME" >/dev/null 2>&1; then
        error "Volume group '$VG_NAME' not found. Are you running on the migrated LVM system?"
    fi
    
    # Check available space for snapshots
    local vg_free=$(vgs --noheadings -o vg_free --units g "$VG_NAME" | tr -d ' G')
    local snapshot_size_num=$(echo "$SNAPSHOT_SIZE" | tr -d 'G')
    local required_space=$((snapshot_size_num * 3))  # 3 snapshots
    
    if (( $(echo "$vg_free < $required_space" | bc -l) )); then
        warning "Low free space in volume group ($vg_free GB available, $required_space GB recommended)"
        echo "Consider reducing snapshot size or extending volume group"
    fi
    
    success "LVM system check passed"
}

detect_backup_drive() {
    log "Detecting external backup drive..."
    
    # Look for drives that are not the current root drive
    local root_drive=$(lsblk -no PKNAME $(findmnt -no SOURCE /))
    local available_drives=($(lsblk -dpno NAME | grep -v "$root_drive"))
    
    if [ ${#available_drives[@]} -eq 0 ]; then
        error "No external drives found for backup"
    elif [ ${#available_drives[@]} -eq 1 ]; then
        EXTERNAL_BACKUP_DRIVE="${available_drives[0]}"
        log "Auto-selected backup drive: $EXTERNAL_BACKUP_DRIVE"
    else
        echo "Multiple external drives found:"
        for i in "${!available_drives[@]}"; do
            local drive="${available_drives[$i]}"
            local info=$(lsblk -no SIZE,MODEL "$drive" | head -1)
            echo "$((i+1)). $drive - $info"
        done
        
        read -p "Select backup drive [1-${#available_drives[@]}]: " choice
        EXTERNAL_BACKUP_DRIVE="${available_drives[$((choice-1))]}"
    fi
    
    success "Selected backup drive: $EXTERNAL_BACKUP_DRIVE"
}

create_snapshots() {
    log "Creating LVM snapshots..."
    
    # Remove existing snapshots if they exist
    remove_snapshots_quiet
    
    # Create snapshots for each logical volume
    local volumes=("root" "home" "boot")
    local snapshots=("$ROOT_SNAPSHOT" "$HOME_SNAPSHOT" "$BOOT_SNAPSHOT")
    
    for i in "${!volumes[@]}"; do
        local volume="${volumes[$i]}"
        local snapshot="${snapshots[$i]}"
        
        if lvs "$VG_NAME/$volume" >/dev/null 2>&1; then
            log "Creating snapshot: $snapshot for $volume"
            lvcreate -L "$SNAPSHOT_SIZE" -s -n "$snapshot" "/dev/$VG_NAME/$volume" || {
                error "Failed to create snapshot $snapshot"
            }
            success "Created snapshot: $snapshot"
        else
            warning "Logical volume $volume not found, skipping"
        fi
    done
    
    success "All snapshots created successfully"
}

mount_snapshots() {
    log "Mounting snapshots for backup..."
    
    # Create mount points
    mkdir -p "$BACKUP_BASE_DIR"/{root,home,boot}
    
    # Mount snapshots
    local snapshots=("$ROOT_SNAPSHOT" "$HOME_SNAPSHOT" "$BOOT_SNAPSHOT")
    local mount_points=("root" "home" "boot")
    
    for i in "${!snapshots[@]}"; do
        local snapshot="${snapshots[$i]}"
        local mount_point="$BACKUP_BASE_DIR/${mount_points[$i]}"
        
        if [ -e "/dev/$VG_NAME/$snapshot" ]; then
            log "Mounting /dev/$VG_NAME/$snapshot to $mount_point"
            mount "/dev/$VG_NAME/$snapshot" "$mount_point" || {
                warning "Failed to mount snapshot $snapshot"
            }
        fi
    done
    
    success "Snapshots mounted at $BACKUP_BASE_DIR"
}

backup_to_external() {
    log "Backing up snapshots to external drive..."
    
    if [ -z "$EXTERNAL_BACKUP_DRIVE" ]; then
        detect_backup_drive
    fi
    
    # Create backup directory on external drive
    local external_mount="/mnt/external-backup"
    mkdir -p "$external_mount"
    
    # Mount external drive (assume first partition)
    mount "${EXTERNAL_BACKUP_DRIVE}1" "$external_mount" 2>/dev/null || {
        # Try the whole drive if partition doesn't work
        mount "$EXTERNAL_BACKUP_DRIVE" "$external_mount" || {
            error "Failed to mount external backup drive"
        }
    }
    
    # Create timestamped backup directory
    local backup_timestamp=$(date '+%Y%m%d_%H%M%S')
    local backup_dir="$external_mount/lvm-snapshots/$backup_timestamp"
    mkdir -p "$backup_dir"
    
    # Backup each mounted snapshot
    local mount_points=("root" "home" "boot")
    for mount_point in "${mount_points[@]}"; do
        local source="$BACKUP_BASE_DIR/$mount_point"
        local target="$backup_dir/$mount_point"
        
        if mountpoint -q "$source"; then
            log "Backing up $mount_point to external drive..."
            mkdir -p "$target"
            rsync -avxHAX --progress "$source/" "$target/" || {
                warning "Backup of $mount_point failed"
            }
        else
            warning "Snapshot $mount_point not mounted, skipping"
        fi
    done
    
    # Create backup metadata
    cat > "$backup_dir/backup-info.txt" << EOF
LVM Snapshot Backup Information
===============================
Backup Date: $(date)
Source System: $(hostname)
Volume Group: $VG_NAME
Snapshot Size: $SNAPSHOT_SIZE
Backup Location: $backup_dir

Logical Volumes:
$(lvs $VG_NAME)

Volume Group Info:
$(vgs $VG_NAME)

Physical Volumes:
$(pvs)
EOF
    
    # Unmount external drive
    umount "$external_mount"
    
    success "Backup completed to $backup_dir"
}

remove_snapshots_quiet() {
    # Remove snapshots without error messages (used internally)
    lvremove -f "/dev/$VG_NAME/$ROOT_SNAPSHOT" 2>/dev/null || true
    lvremove -f "/dev/$VG_NAME/$HOME_SNAPSHOT" 2>/dev/null || true  
    lvremove -f "/dev/$VG_NAME/$BOOT_SNAPSHOT" 2>/dev/null || true
}

remove_snapshots() {
    log "Cleaning up snapshots..."
    
    # Unmount snapshots
    umount "$BACKUP_BASE_DIR/root" 2>/dev/null || true
    umount "$BACKUP_BASE_DIR/home" 2>/dev/null || true
    umount "$BACKUP_BASE_DIR/boot" 2>/dev/null || true
    
    # Remove logical volumes
    remove_snapshots_quiet
    
    success "Snapshots removed"
}

list_snapshots() {
    echo "Current LVM snapshots:"
    lvs "$VG_NAME" | grep -E "(snapshot|Snap%)" || echo "No snapshots found"
}

show_status() {
    echo -e "${GREEN}=== LVM Snapshot Backup Status ===${NC}"
    echo
    
    echo "Volume Group Information:"
    vgs "$VG_NAME" 2>/dev/null || echo "Volume group not found"
    echo
    
    echo "Logical Volumes:"
    lvs "$VG_NAME" 2>/dev/null || echo "No logical volumes found"
    echo
    
    echo "Current Snapshots:"
    list_snapshots
    echo
    
    echo "Mount Status:"
    if mountpoint -q "$BACKUP_BASE_DIR/root" 2>/dev/null; then
        echo "  Root snapshot: mounted at $BACKUP_BASE_DIR/root"
    else
        echo "  Root snapshot: not mounted"
    fi
    
    if mountpoint -q "$BACKUP_BASE_DIR/home" 2>/dev/null; then
        echo "  Home snapshot: mounted at $BACKUP_BASE_DIR/home"
    else
        echo "  Home snapshot: not mounted"
    fi
    
    if mountpoint -q "$BACKUP_BASE_DIR/boot" 2>/dev/null; then
        echo "  Boot snapshot: mounted at $BACKUP_BASE_DIR/boot"
    else
        echo "  Boot snapshot: not mounted"
    fi
}

integrate_with_backup_system() {
    log "Integrating with existing backup system..."
    
    # Check if main backup system exists
    if [ -f "./backup_script.sh" ]; then
        log "Found main backup system - creating integration"
        
        # Create wrapper script that combines snapshot and full backup
        cat > "./lvm_integrated_backup.sh" << 'EOF'
#!/bin/bash

# Integrated LVM + Full System Backup
# Combines LVM snapshots with the main backup system

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

echo "=== Integrated LVM Backup System ==="
echo "1. LVM Snapshot Backup (fast, consistent)"
echo "2. Full System Backup (complete clone)"
echo "3. Combined Backup (snapshots + clone)"
echo

read -p "Select backup type [1-3]: " choice

case $choice in
    1)
        echo "Performing LVM snapshot backup..."
        sudo "$SCRIPT_DIR/lvm_snapshot_backup.sh" backup
        ;;
    2)
        echo "Performing full system backup..."
        sudo "$SCRIPT_DIR/backup_script.sh"
        ;;
    3)
        echo "Performing combined backup..."
        echo "Step 1: LVM snapshots..."
        sudo "$SCRIPT_DIR/lvm_snapshot_backup.sh" backup
        echo "Step 2: Full system clone..."
        sudo "$SCRIPT_DIR/backup_script.sh" --sync
        ;;
    *)
        echo "Invalid choice"
        exit 1
        ;;
esac
EOF
        chmod +x "./lvm_integrated_backup.sh"
        success "Created integrated backup script: lvm_integrated_backup.sh"
    else
        warning "Main backup system not found in current directory"
    fi
}

show_usage() {
    echo "Usage: $0 {create|mount|remove|backup|external|status|list|integrate}"
    echo
    echo "Commands:"
    echo "  create     - Create LVM snapshots only"
    echo "  mount      - Mount existing snapshots for access"
    echo "  remove     - Remove snapshots and unmount"
    echo "  backup     - Create snapshots and mount (ready for backup)"
    echo "  external   - Full backup to external drive"
    echo "  status     - Show current LVM and snapshot status"
    echo "  list       - List current snapshots"
    echo "  integrate  - Integrate with existing backup system"
    echo
    echo "Examples:"
    echo "  $0 backup           # Create and mount snapshots"
    echo "  $0 external         # Full backup to external drive"
    echo "  $0 remove           # Clean up when backup complete"
}

main() {
    # Ensure log file exists and is writable
    sudo touch "$LOG_FILE" 2>/dev/null || LOG_FILE="/tmp/lvm-backup.log"
    
    case "${1:-}" in
        create)
            check_lvm_system
            create_snapshots
            ;;
        mount)
            mount_snapshots
            ;;
        remove)
            remove_snapshots
            ;;
        backup)
            check_lvm_system
            create_snapshots
            mount_snapshots
            echo "Snapshots ready for backup at $BACKUP_BASE_DIR"
            echo "Run '$0 remove' when backup is complete"
            ;;
        external)
            check_lvm_system
            create_snapshots
            mount_snapshots
            backup_to_external
            remove_snapshots
            ;;
        status)
            show_status
            ;;
        list)
            list_snapshots
            ;;
        integrate)
            integrate_with_backup_system
            ;;
        *)
            show_usage
            exit 1
            ;;
    esac
}

# Check if running as root for LVM operations
if [ "$EUID" -ne 0 ] && [[ "$1" =~ ^(create|mount|remove|backup|external)$ ]]; then
    error "LVM operations require root privileges. Use: sudo $0 $1"
fi

main "$@"