#!/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"
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
}

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, unmounting first..." | tee -a "$LOG"
        umount "$MOUNT_POINT"
    fi
    
    # 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"
    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
        umount "$MOUNT_POINT"
    fi
}

# === 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 streaming 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 one by one
    for chunk in "$LATEST_BACKUP_DIR"/chunk_*; 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/"
        
        # If this is the first chunk, start extraction
        if [ "$(basename "$chunk")" = "chunk_aa" ]; then
            7z x -y "$TMP_RESTORE_DIR/$(basename "$chunk")" -o/ &
            extract_pid=$!
        else
            # For subsequent chunks, append to the extraction
            cat "$TMP_RESTORE_DIR/$(basename "$chunk")" >> "$TMP_RESTORE_DIR/combined.7z"
        fi
        
        # Clean up the processed chunk
        rm -f "$TMP_RESTORE_DIR/$(basename "$chunk")"
    done
    
    # Wait for extraction to complete
    wait $extract_pid
    
    # 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"

# Create a temporary file list
echo "[$(date)] 📋 Creating file list..." | tee -a "$LOG"
find $SRC -type f > "$TMP_BACKUP/filelist.txt"

# Process files in chunks
while IFS= read -r file; do
    # Check space before processing each file
    if ! check_space; then
        echo "[$(date)] ⏳ Waiting for more space..." | tee -a "$LOG"
        wait_for_space
    fi
    
    # Get the relative path for the file
    rel_path="${file#/}"
    chunk_dir="$TMP_BACKUP/$(dirname "$rel_path")"
    mkdir -p "$chunk_dir"
    
    # Compress the file
    echo "[$(date)] 📦 Compressing: $file" | tee -a "$LOG"
    7z a -t7z -m0=lzma2 -mx=5 "$TMP_BACKUP/$(echo "$rel_path" | tr / _).7z" "$file"
    
    # If we have enough chunks or space is low, transfer them
    if [ $(ls "$TMP_BACKUP"/*.7z 2>/dev/null | wc -l) -ge 5 ] || ! check_space; then
        echo "[$(date)] ⬆️ Transferring chunks to SMB..." | tee -a "$LOG"
        for chunk in "$TMP_BACKUP"/*.7z; do
            if [ -f "$chunk" ]; then
                rsync -ah --progress "$chunk" "$BACKUP_DIR/"
                rm -f "$chunk"
            fi
        done
    fi
done < "$TMP_BACKUP/filelist.txt"

# Transfer any remaining chunks
echo "[$(date)] ⬆️ Transferring remaining chunks..." | tee -a "$LOG"
for chunk in "$TMP_BACKUP"/*.7z; do
    if [ -f "$chunk" ]; then
        rsync -ah --progress "$chunk" "$BACKUP_DIR/"
        rm -f "$chunk"
    fi
done

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