#!/bin/bash

# === Config ===
SERVER_IP="192.168.2.194"
SHARE_NAME="PBS-back"
MOUNT_POINT="/mnt/smb-backup"
CREDENTIALS_FILE="/root/.smb-pbs-cred"
DATE=$(date +%F-%H%M)
BACKUP_NAME="pbs-config-$DATE"
TMP_BACKUP="/tmp/$BACKUP_NAME"
CHUNK_SIZE="10G"
MIN_SPACE_REQUIRED="5G"  # Minimum space required in /tmp
LOG="/var/log/pbs-smb-backup.log"
SRC="/etc /root /mypool"

# Exclusion patterns
EXCLUDE_PATTERNS=(
    "node_modules"
    ".cursor-server"
    ".git"
    ".vscode"
    "*.log"
    "*.tmp"
    "*.temp"
    "*.swp"
    "*.swo"
    "*.bak"
    "*.cache"
    "*.pid"
    "*.lock"
    "*.socket"
    "*.sock"
    "*.pid"
    "*.log.*"
    "*.gz"
    "*.zip"
    "*.tar"
    "*.7z"
    "*.rar"
    "*.iso"
    "*.img"
    "*.vmdk"
    "*.vhd"
    "*.qcow2"
    "*.raw"
    "*.vdi"
    "*.vbox"
    "*.ova"
    "*.ovf"
    "*.vmx"
    "*.vmsd"
    "*.vmsn"
    "*.vmss"
    "*.vmtm"
    "*.vmxf"
    "*.nvram"
    "*.vmem"
    "*.vswp"
)

MAX_BACKUPS=3  # Keep exactly 3 backups

# === Cleanup function ===
cleanup_old_backups() {
    echo "[$(date)] 🧹 Starting cleanup of old backups..." | tee -a "$LOG"
    
    # Cleanup old backups on SMB share
    if mountpoint -q "$MOUNT_POINT"; then
        echo "[$(date)] Keeping only the $MAX_BACKUPS most recent backups..." | tee -a "$LOG"
        # List all backup directories, sort by date (newest first), and remove older ones
        cd "$MOUNT_POINT" || exit 1
        ls -td pbs-config-* 2>/dev/null | tail -n +$((MAX_BACKUPS + 1)) | xargs -r rm -rf
        echo "[$(date)] Current backups:" | tee -a "$LOG"
        ls -lhd pbs-config-* 2>/dev/null | tee -a "$LOG"
    fi
    
    # Cleanup old logs
    echo "[$(date)] Cleaning up old log files..." | tee -a "$LOG"
    find /var/log -name "pbs-smb-backup.*.log" -mtime +7 -delete
    
    # Cleanup /tmp
    echo "[$(date)] Cleaning up old temporary files..." | tee -a "$LOG"
    rm -rf /tmp/pbs-config-*
    rm -rf /tmp/restore-*
}

# === Create mount point if needed ===
mkdir -p "$MOUNT_POINT"

# === Ensure credentials file exists ===
if [ ! -f "$CREDENTIALS_FILE" ]; then
    echo "username=pbs" > "$CREDENTIALS_FILE"
    echo "password=2104" >> "$CREDENTIALS_FILE"
    chmod 600 "$CREDENTIALS_FILE"
fi

# === Helper functions ===
get_available_space() {
    df -B1 /tmp | awk 'NR==2 {print $4}'
}

check_space() {
    local required_space=$(numfmt --from=iec $MIN_SPACE_REQUIRED)
    local available_space=$(get_available_space)
    
    if [ "$available_space" -lt "$required_space" ]; then
        echo "[$(date)] ⚠️ Warning: Only $(numfmt --to=iec $available_space) available in /tmp" | tee -a "$LOG"
        return 1
    fi
    return 0
}

wait_for_space() {
    local required_space=$(numfmt --from=iec $MIN_SPACE_REQUIRED)
    local available_space
    
    while true; do
        available_space=$(get_available_space)
        if [ "$available_space" -ge "$required_space" ]; then
            return 0
        fi
        echo "[$(date)] ⚠️ Waiting for more space in /tmp (need $MIN_SPACE_REQUIRED, have $(numfmt --to=iec $available_space))..." | tee -a "$LOG"
        sleep 10
    done
}

cleanup_processes() {
    local mount_point="$1"
    local max_attempts=3
    local attempt=1
    
    echo "[$(date)] Cleaning up processes using $mount_point..." | tee -a "$LOG"
    
    while [ $attempt -le $max_attempts ]; do
        # Get list of processes using the mount point, excluding our own script
        local pids=$(lsof -t "$mount_point" 2>/dev/null | grep -v "$$")
        if [ -z "$pids" ]; then
            echo "[$(date)] No processes found using $mount_point" | tee -a "$LOG"
            return 0
        fi
        
        echo "[$(date)] Attempt $attempt: Found processes: $pids" | tee -a "$LOG"
        
        # First try SIGTERM
        kill -15 $pids 2>/dev/null
        sleep 5
        
        # Check if processes are still running
        pids=$(lsof -t "$mount_point" 2>/dev/null | grep -v "$$")
        if [ -n "$pids" ]; then
            echo "[$(date)] Processes still running, sending SIGKILL..." | tee -a "$LOG"
            kill -9 $pids 2>/dev/null
            sleep 5
        fi
        
        # Final check
        if [ -z "$(lsof -t "$mount_point" 2>/dev/null | grep -v "$$")" ]; then
            echo "[$(date)] Successfully cleaned up all processes" | tee -a "$LOG"
            return 0
        fi
        
        ((attempt++))
    done
    
    echo "[$(date)] ⚠️ Warning: Could not clean up all processes after $max_attempts attempts" | tee -a "$LOG"
    return 1
}

mount_share() {
    echo "[$(date)] Mounting SMB share..." | tee -a "$LOG"
    
    # Check if already mounted
    if mountpoint -q "$MOUNT_POINT"; then
        echo "[$(date)] Share already mounted, attempting to unmount..." | tee -a "$LOG"
        
        # Clean up any processes first
        cleanup_processes "$MOUNT_POINT"
        
        # Try to unmount gracefully first
        umount "$MOUNT_POINT" 2>/dev/null
        sleep 2
        
        # If still mounted, try force unmount
        if mountpoint -q "$MOUNT_POINT"; then
            echo "[$(date)] Force unmounting..." | tee -a "$LOG"
            umount -f "$MOUNT_POINT" 2>/dev/null
            sleep 2
            
            # If still mounted, try lazy unmount
            if mountpoint -q "$MOUNT_POINT"; then
                echo "[$(date)] Lazy unmounting..." | tee -a "$LOG"
                umount -l "$MOUNT_POINT" 2>/dev/null
                sleep 2
                
                # If still mounted, try one more time with process cleanup
                if mountpoint -q "$MOUNT_POINT"; then
                    cleanup_processes "$MOUNT_POINT"
                    sleep 2
                    umount -f "$MOUNT_POINT" 2>/dev/null
                    sleep 2
                fi
            fi
        fi
        
        # Final check - if still mounted, exit
        if mountpoint -q "$MOUNT_POINT"; then
            echo "[$(date)] ❌ Failed to unmount existing share. Please check for processes using $MOUNT_POINT" | tee -a "$LOG"
            echo "[$(date)] Running processes:" | tee -a "$LOG"
            lsof "$MOUNT_POINT" | grep -v "$$" | tee -a "$LOG"
            exit 1
        fi
    fi
    
    # Ensure mount point is clean
    rm -rf "$MOUNT_POINT"/*
    
    # Try mounting with different SMB versions if needed
    for smb_ver in "3.0" "2.1" "2.0" "1.0"; do
        echo "[$(date)] Attempting mount with SMB version $smb_ver..." | tee -a "$LOG"
        mount -t cifs "//$SERVER_IP/$SHARE_NAME" "$MOUNT_POINT" -o "credentials=$CREDENTIALS_FILE,iocharset=utf8,vers=$smb_ver,sec=ntlmssp,uid=0,gid=0,file_mode=0644,dir_mode=0755"
        if [ $? -eq 0 ]; then
            echo "[$(date)] Successfully mounted with SMB version $smb_ver" | tee -a "$LOG"
            return 0
        fi
        # Show dmesg output for debugging
        echo "[$(date)] Mount failed, checking dmesg for details..." | tee -a "$LOG"
        dmesg | tail -n 20 | tee -a "$LOG"
        sleep 2
    done
    
    echo "[$(date)] ❌ Failed to mount SMB share. Please check:" | tee -a "$LOG"
    echo "1. SMB server is running and accessible" | tee -a "$LOG"
    echo "2. Credentials are correct" | tee -a "$LOG"
    echo "3. Network connectivity to $SERVER_IP" | tee -a "$LOG"
    echo "4. SMB share '$SHARE_NAME' exists and is accessible" | tee -a "$LOG"
    exit 1
}

unmount_share() {
    echo "[$(date)] Unmounting SMB share..." | tee -a "$LOG"
    if mountpoint -q "$MOUNT_POINT"; then
        # Clean up any processes first
        cleanup_processes "$MOUNT_POINT"
        
        # Try to unmount gracefully first
        umount "$MOUNT_POINT" 2>/dev/null
        sleep 2
        
        # If still mounted, try force unmount
        if mountpoint -q "$MOUNT_POINT"; then
            echo "[$(date)] Force unmounting..." | tee -a "$LOG"
            umount -f "$MOUNT_POINT" 2>/dev/null
            sleep 2
            
            # If still mounted, try lazy unmount
            if mountpoint -q "$MOUNT_POINT"; then
                echo "[$(date)] Lazy unmounting..." | tee -a "$LOG"
                umount -l "$MOUNT_POINT" 2>/dev/null
                sleep 2
                
                # If still mounted, try one more time with process cleanup
                if mountpoint -q "$MOUNT_POINT"; then
                    cleanup_processes "$MOUNT_POINT"
                    sleep 2
                    umount -f "$MOUNT_POINT" 2>/dev/null
                    sleep 2
                fi
            fi
        fi
        
        # Final check - if still mounted, warn but continue
        if mountpoint -q "$MOUNT_POINT"; then
            echo "[$(date)] ⚠️ Warning: Could not unmount $MOUNT_POINT completely" | tee -a "$LOG"
            echo "[$(date)] Running processes:" | tee -a "$LOG"
            lsof "$MOUNT_POINT" | grep -v "$$" | tee -a "$LOG"
        fi
    fi
}

# Add trap to ensure cleanup on script exit
trap 'handle_exit' EXIT

handle_exit() {
    local exit_code=$?
    echo "[$(date)] Cleaning up on exit..." | tee -a "$LOG"
    
    # Only attempt cleanup if we're not already in a cleanup process
    if [ -z "$CLEANUP_IN_PROGRESS" ]; then
        export CLEANUP_IN_PROGRESS=1
        cleanup_processes "$MOUNT_POINT"
        unmount_share
        rm -rf "$TMP_BACKUP"
    fi
    
    exit $exit_code
}

# === Handle -u option for restore ===
if [ "$1" == "-u" ]; then
    echo "[$(date)] 🔁 Restore mode selected. Starting cleanup..." | tee -a "$LOG"
    cleanup_old_backups
    mount_share

    # Get latest backup directory
    LATEST_BACKUP_DIR=$(ls -td "$MOUNT_POINT"/pbs-config-* 2>/dev/null | head -n1)

    if [ -z "$LATEST_BACKUP_DIR" ]; then
        echo "[$(date)] ❌ No backup found on SMB share." | tee -a "$LOG"
        unmount_share
        exit 1
    fi

    echo "[$(date)] ⬇️ Starting restore process from $LATEST_BACKUP_DIR..." | tee -a "$LOG"
    
    # Create a temporary directory for processing
    TMP_RESTORE_DIR="/tmp/restore-$(basename "$LATEST_BACKUP_DIR")"
    mkdir -p "$TMP_RESTORE_DIR"
    
    # Process chunks in order
    for chunk in $(ls -v "$LATEST_BACKUP_DIR"/chunk_* 2>/dev/null); do
        if [ ! -f "$chunk" ]; then
            continue
        fi
        
        echo "[$(date)] Processing chunk: $(basename "$chunk")" | tee -a "$LOG"
        
        # Copy chunk to temp
        rsync -ah --progress "$chunk" "$TMP_RESTORE_DIR/"
        
        # Extract the chunk
        echo "[$(date)] Extracting chunk..." | tee -a "$LOG"
        7z x -y "$TMP_RESTORE_DIR/$(basename "$chunk")" -o/ | tee -a "$LOG"
        
        # Remove the processed chunk
        rm -f "$TMP_RESTORE_DIR/$(basename "$chunk")"
    done
    
    # Cleanup
    rm -rf "$TMP_RESTORE_DIR"
    unmount_share
    echo "[$(date)] ✅ Restore completed." | tee -a "$LOG"
    exit 0
fi

# === Regular backup mode ===

# Step 1: Start cleanup
echo "[$(date)] 🧹 Starting initial cleanup..." | tee -a "$LOG"
cleanup_old_backups

if ! check_space; then
    echo "[$(date)] ⚠️ Proceeding with caution due to limited space" | tee -a "$LOG"
fi

# Step 2: Create temporary directory for chunks
mkdir -p "$TMP_BACKUP"

# Step 3: Mount share
mount_share

# Step 4: Create backup directory on SMB
BACKUP_DIR="$MOUNT_POINT/$BACKUP_NAME"
mkdir -p "$BACKUP_DIR"

# Step 5: Create backup in chunks with space management
echo "[$(date)] 🗜️ Starting backup with space management..." | tee -a "$LOG"

# Build find command exclusions
build_find_exclusions() {
    local exclusions=""
    for pattern in "${EXCLUDE_PATTERNS[@]}"; do
        exclusions+=" -not -path '*/$pattern/*' -not -name '$pattern'"
    done
    echo "$exclusions"
}

# Create a temporary file list with exclusions
echo "[$(date)] 📋 Creating file list (excluding unnecessary files)..." | tee -a "$LOG"
eval "find $SRC -type f $(build_find_exclusions) > \"$TMP_BACKUP/filelist.txt\""

# Initialize variables for chunk creation
chunk_num=1
current_chunk_size=0
current_chunk_files=()

# Process files and create chunks
while IFS= read -r file; do
    # Skip if file doesn't exist
    [ -f "$file" ] || continue
    
    # Get file size
    file_size=$(stat -c %s "$file")
    
    # If adding this file would exceed chunk size, create the chunk
    if [ $((current_chunk_size + file_size)) -gt $(numfmt --from=iec $CHUNK_SIZE) ]; then
        if [ ${#current_chunk_files[@]} -gt 0 ]; then
            echo "[$(date)] 📦 Creating chunk $chunk_num with ${#current_chunk_files[@]} files..." | tee -a "$LOG"
            
            # Create the chunk and wait for it to complete
            chunk_file="$TMP_BACKUP/chunk_$(printf "%03d" $chunk_num).7z"
            7z a -y -spf -t7z -m0=lzma2 -mx=5 "$chunk_file" "${current_chunk_files[@]}" | tee -a "$LOG"
            
            # Wait for 7z to complete and verify the chunk exists
            if [ -f "$chunk_file" ]; then
                echo "[$(date)] ⬆️ Transferring chunk $chunk_num to SMB..." | tee -a "$LOG"
                rsync -ah --progress "$chunk_file" "$BACKUP_DIR/"
                
                # Verify the transfer was successful
                if [ $? -eq 0 ]; then
                    echo "[$(date)] ✅ Chunk $chunk_num transferred successfully" | tee -a "$LOG"
                    rm -f "$chunk_file"
                else
                    echo "[$(date)] ❌ Failed to transfer chunk $chunk_num" | tee -a "$LOG"
                    exit 1
                fi
            else
                echo "[$(date)] ❌ Failed to create chunk $chunk_num" | tee -a "$LOG"
                exit 1
            fi
            
            # Reset for next chunk
            current_chunk_files=()
            current_chunk_size=0
            ((chunk_num++))
            
            # Check space before continuing
            if ! check_space; then
                wait_for_space
            fi
        fi
    fi
    
    # Add file to current chunk
    current_chunk_files+=("$file")
    current_chunk_size=$((current_chunk_size + file_size))
done < "$TMP_BACKUP/filelist.txt"

# Create final chunk if there are remaining files
if [ ${#current_chunk_files[@]} -gt 0 ]; then
    echo "[$(date)] 📦 Creating final chunk $chunk_num with ${#current_chunk_files[@]} files..." | tee -a "$LOG"
    chunk_file="$TMP_BACKUP/chunk_$(printf "%03d" $chunk_num).7z"
    7z a -y -spf -t7z -m0=lzma2 -mx=5 "$chunk_file" "${current_chunk_files[@]}" | tee -a "$LOG"
    
    # Wait for 7z to complete and verify the chunk exists
    if [ -f "$chunk_file" ]; then
        echo "[$(date)] ⬆️ Transferring final chunk to SMB..." | tee -a "$LOG"
        rsync -ah --progress "$chunk_file" "$BACKUP_DIR/"
        
        # Verify the transfer was successful
        if [ $? -eq 0 ]; then
            echo "[$(date)] ✅ Final chunk transferred successfully" | tee -a "$LOG"
            rm -f "$chunk_file"
        else
            echo "[$(date)] ❌ Failed to transfer final chunk" | tee -a "$LOG"
            exit 1
        fi
    else
        echo "[$(date)] ❌ Failed to create final chunk" | tee -a "$LOG"
        exit 1
    fi
fi

# Step 6: Final cleanup and finish
rm -rf "$TMP_BACKUP"
unmount_share
echo "[$(date)] ✅ Backup completed successfully: $BACKUP_NAME" | tee -a "$LOG"