#!/bin/bash

# Enhanced Simple LVM Backup Script
# Supports three backup modes:
# 1. LV → LV: Update existing logical volume backup
# 2. LV → Raw: Create fresh backup on raw device  
# 3. VG → Raw: Clone entire volume group to raw device

set -e

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

usage() {
    echo "Enhanced Simple LVM Backup Script"
    echo ""
    echo "Usage:"
    echo "  $0 lv-to-lv SOURCE_LV TARGET_LV"
    echo "  $0 lv-to-raw SOURCE_LV TARGET_DEVICE"
    echo "  $0 vg-to-raw SOURCE_VG TARGET_DEVICE"
    echo ""
    echo "Modes:"
    echo "  lv-to-lv   - Update existing LV backup (SOURCE_LV → TARGET_LV)"
    echo "  lv-to-raw  - Create fresh backup (SOURCE_LV → raw device)"
    echo "  vg-to-raw  - Clone entire VG (SOURCE_VG → raw device)"
    echo ""
    echo "Examples:"
    echo "  $0 lv-to-lv /dev/internal-vg/root /dev/backup-vg/root"
    echo "  $0 lv-to-raw /dev/internal-vg/root /dev/sdb"
    echo "  $0 vg-to-raw internal-vg /dev/sdb"
    echo ""
    echo "List available sources/targets:"
    echo "  ./list_drives.sh"
    echo ""
    exit 1
}

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

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

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

info() {
    echo -e "${BLUE}[INFO] $1${NC}"
}

cleanup_and_exit() {
    local exit_code=${1:-0}
    
    if [ -n "$SNAPSHOT_PATH" ] && lvs "$SNAPSHOT_PATH" >/dev/null 2>&1; then
        warn "Cleaning up snapshot: $SNAPSHOT_PATH"
        lvremove -f "$SNAPSHOT_PATH" || warn "Failed to remove snapshot"
    fi
    
    exit $exit_code
}

# Trap for cleanup
trap 'cleanup_and_exit 130' INT TERM

# Check arguments
if [ $# -ne 3 ]; then
    usage
fi

MODE="$1"
SOURCE="$2"
TARGET="$3"

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    error "This script must be run as root"
fi

# Validate mode
case "$MODE" in
    "lv-to-lv"|"lv-to-raw"|"vg-to-raw")
        ;;
    *)
        error "Invalid mode: $MODE"
        ;;
esac

log "Enhanced Simple LVM Backup"
log "Mode: $MODE"
log "Source: $SOURCE"
log "Target: $TARGET"

# Mode-specific validation and execution
case "$MODE" in
    "lv-to-lv")
        # LV to LV backup
        if [ ! -e "$SOURCE" ]; then
            error "Source LV does not exist: $SOURCE"
        fi
        
        if [ ! -e "$TARGET" ]; then
            error "Target LV does not exist: $TARGET"
        fi
        
        # Extract VG and LV names
        VG_NAME=$(lvs --noheadings -o vg_name "$SOURCE" | tr -d ' ')
        LV_NAME=$(lvs --noheadings -o lv_name "$SOURCE" | tr -d ' ')
        SNAPSHOT_NAME="${LV_NAME}_backup_snap"
        SNAPSHOT_PATH="/dev/$VG_NAME/$SNAPSHOT_NAME"
        
        info "This will update the existing backup LV"
        echo ""
        echo -e "${YELLOW}Source LV: $SOURCE${NC}"
        echo -e "${YELLOW}Target LV: $TARGET${NC}"
        echo -e "${YELLOW}The target LV will be overwritten with current source data${NC}"
        echo ""
        read -p "Continue? (yes/no): " confirm
        
        if [ "$confirm" != "yes" ]; then
            echo "Backup cancelled."
            exit 0
        fi
        
        log "Creating snapshot of source LV"
        lvcreate -L1G -s -n "$SNAPSHOT_NAME" "$SOURCE" || error "Failed to create snapshot"
        
        log "Copying snapshot to target LV"
        if command -v pv >/dev/null 2>&1; then
            pv "$SNAPSHOT_PATH" | dd of="$TARGET" bs=4M || error "Copy failed"
        else
            dd if="$SNAPSHOT_PATH" of="$TARGET" bs=4M status=progress || error "Copy failed"
        fi
        
        log "Cleaning up snapshot"
        lvremove -f "$SNAPSHOT_PATH" || warn "Failed to remove snapshot"
        SNAPSHOT_PATH=""
        
        log "LV to LV backup completed successfully"
        ;;
        
    "lv-to-raw")
        # LV to raw device backup
        if [ ! -e "$SOURCE" ]; then
            error "Source LV does not exist: $SOURCE"
        fi
        
        if [ ! -e "$TARGET" ]; then
            error "Target device does not exist: $TARGET"
        fi
        
        # Extract VG and LV names
        VG_NAME=$(lvs --noheadings -o vg_name "$SOURCE" | tr -d ' ')
        LV_NAME=$(lvs --noheadings -o lv_name "$SOURCE" | tr -d ' ')
        SNAPSHOT_NAME="${LV_NAME}_backup_snap"
        SNAPSHOT_PATH="/dev/$VG_NAME/$SNAPSHOT_NAME"
        
        info "This will create a fresh backup on raw device"
        echo ""
        echo -e "${YELLOW}Source LV: $SOURCE${NC}"
        echo -e "${YELLOW}Target Device: $TARGET${NC}"
        echo -e "${RED}WARNING: All data on $TARGET will be lost!${NC}"
        echo ""
        read -p "Continue? (yes/no): " confirm
        
        if [ "$confirm" != "yes" ]; then
            echo "Backup cancelled."
            exit 0
        fi
        
        log "Creating snapshot of source LV"
        lvcreate -L1G -s -n "$SNAPSHOT_NAME" "$SOURCE" || error "Failed to create snapshot"
        
        log "Copying snapshot to target device"
        if command -v pv >/dev/null 2>&1; then
            pv "$SNAPSHOT_PATH" | dd of="$TARGET" bs=4M || error "Copy failed"
        else
            dd if="$SNAPSHOT_PATH" of="$TARGET" bs=4M status=progress || error "Copy failed"
        fi
        
        log "Cleaning up snapshot"
        lvremove -f "$SNAPSHOT_PATH" || warn "Failed to remove snapshot"
        SNAPSHOT_PATH=""
        
        log "LV to raw device backup completed successfully"
        ;;
        
    "vg-to-raw")
        # VG to raw device backup
        if ! vgs "$SOURCE" >/dev/null 2>&1; then
            error "Source VG does not exist: $SOURCE"
        fi
        
        if [ ! -e "$TARGET" ]; then
            error "Target device does not exist: $TARGET"
        fi
        
        # Get the first PV of the source VG
        SOURCE_PV=$(vgs --noheadings -o pv_name "$SOURCE" | head -n1 | tr -d ' ')
        
        if [ -z "$SOURCE_PV" ]; then
            error "No physical volumes found for VG: $SOURCE"
        fi
        
        info "This will clone the entire volume group"
        echo ""
        echo -e "${YELLOW}Source VG: $SOURCE${NC}"
        echo -e "${YELLOW}Source PV: $SOURCE_PV${NC}"
        echo -e "${YELLOW}Target Device: $TARGET${NC}"
        echo -e "${RED}WARNING: All data on $TARGET will be lost!${NC}"
        echo -e "${BLUE}This preserves LVM metadata and all logical volumes${NC}"
        echo ""
        read -p "Continue? (yes/no): " confirm
        
        if [ "$confirm" != "yes" ]; then
            echo "Backup cancelled."
            exit 0
        fi
        
        log "Copying entire PV to target device"
        log "This preserves LVM structure and all LVs"
        
        if command -v pv >/dev/null 2>&1; then
            pv "$SOURCE_PV" | dd of="$TARGET" bs=4M || error "Copy failed"
        else
            dd if="$SOURCE_PV" of="$TARGET" bs=4M status=progress || error "Copy failed"
        fi
        
        log "VG to raw device backup completed successfully"
        log "Target device now contains complete LVM structure"
        ;;
esac

echo ""
echo -e "${GREEN}SUCCESS: Backup completed!${NC}"
echo ""
echo "Next steps:"
case "$MODE" in
    "lv-to-lv")
        echo "- Your backup LV has been updated"
        echo "- You can mount $TARGET to verify the backup"
        ;;
    "lv-to-raw")
        echo "- Your backup device contains a raw copy of the LV"
        echo "- You can mount $TARGET directly or restore from it"
        ;;
    "vg-to-raw")
        echo "- Your backup device contains the complete LVM structure"
        echo "- You can boot from $TARGET or import the VG for recovery"
        echo "- To access: vgimport or boot directly from the device"
        ;;
esac
echo ""