#!/bin/bash
#
# update_snapapi.sh - Update/rebuild SnapAPI kernel module
#
# This script can:
#   - Install a custom SnapAPI RPM and rebuild the kernel module
#   - Replace snapapi26.c source file and rebuild
#   - Simply rebuild the current version
#
# Usage:
#   sudo ./update_snapapi.sh --rpm /path/to/snapapi26_modules-X.Y.Z.rpm
#   sudo ./update_snapapi.sh --file /path/to/snapapi26.c
#   sudo ./update_snapapi.sh --rebuild
#
# Based on Acronis KB documentation:
#   https://care.acronis.com/s/article/Acronis-backup-software-Using-custom-SnapAPI-Module-in-Linux
#

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Global variables
MODE=""
RPM_FILE=""
SOURCE_FILE=""
NEW_VERSION=""
CURRENT_VERSION=""
USE_ACRONIS_RPMDB=false
DKMS_SOURCE_DIR="/usr/src/snapapi26"

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

usage() {
    echo "Usage: $0 <mode> [arguments]"
    echo ""
    echo "Modes:"
    echo "  --rpm <file>     Install SnapAPI from custom RPM and rebuild"
    echo "  --file <file>    Replace snapapi26.c with specified file and rebuild"
    echo "  --rebuild        Rebuild current SnapAPI version"
    echo ""
    echo "Examples:"
    echo "  sudo $0 --rpm /tmp/snapapi26_modules-0.8.18-1.noarch.rpm"
    echo "  sudo $0 --file /tmp/snapapi26.c"
    echo "  sudo $0 --rebuild"
    echo ""
    exit 1
}

check_root() {
    if [ "$(id -u)" -ne 0 ]; then
        log_error "This script must be run as root (use sudo)"
        exit 1
    fi
}

detect_system() {
    # If dpkg exists, this is a Debian-based system - use Acronis RPM database
    # Otherwise, it's a native RPM system - use system RPM database
    if command -v dpkg &>/dev/null; then
        USE_ACRONIS_RPMDB=true
        log_info "Debian-based system detected, will use Acronis RPM database"
    else
        USE_ACRONIS_RPMDB=false
        log_info "RPM-based system detected, will use system RPM database"
    fi
}

check_prerequisites() {
    if ! command -v rpm &>/dev/null; then
        log_error "rpm is not installed"
        if [ "$USE_ACRONIS_RPMDB" = true ]; then
            log_error "Install it with: apt-get install rpm"
        fi
        exit 1
    fi

    if ! command -v dkms &>/dev/null; then
        log_error "dkms is not installed"
        if [ "$USE_ACRONIS_RPMDB" = true ]; then
            log_error "Install it with: apt-get install dkms"
        else
            log_error "Install it with your package manager (dnf/yum/zypper install dkms)"
        fi
        exit 1
    fi
}

get_current_snapapi_version() {
    # Get currently installed snapapi version from dkms
    CURRENT_VERSION=$(dkms status snapapi26 2>/dev/null | head -1 | sed -n 's/snapapi26[,/]\s*\([0-9.]*\).*/\1/p' || echo "")
    
    # If not in DKMS, find highest version in /usr/src/snapapi26-*
    if [ -z "$CURRENT_VERSION" ]; then
        CURRENT_VERSION=$(ls -d /usr/src/snapapi26-* 2>/dev/null | sed 's|.*/snapapi26-||' | sort -V | tail -1 || echo "")
    fi
    
    if [ -n "$CURRENT_VERSION" ]; then
        log_info "Current SnapAPI version: $CURRENT_VERSION"
    else
        log_info "No SnapAPI version found"
    fi
}

extract_rpm_version() {
    local rpm_file="$1"
    # Extract version from RPM filename or query RPM
    if command -v rpm &>/dev/null; then
        NEW_VERSION=$(rpm -qp --queryformat '%{VERSION}' "$rpm_file" 2>/dev/null || echo "")
    fi
    
    if [ -z "$NEW_VERSION" ]; then
        # Fallback: extract from filename
        NEW_VERSION=$(basename "$rpm_file" | sed -n 's/snapapi26[_-]modules[_-]\([0-9.]*\).*/\1/p')
    fi
    
    if [ -z "$NEW_VERSION" ]; then
        log_error "Could not determine SnapAPI version from RPM"
        exit 1
    fi
    
    log_info "New SnapAPI version to install: $NEW_VERSION"
}

stop_acronis_services() {
    log_info "Stopping Acronis services..."
    
    # Try systemd first
    for svc in acronis_mms acronis_agent acronis_scheduler; do
        systemctl stop "$svc" 2>/dev/null || true
    done
    
    # Fallback to service command
    for svc in acronis_mms acronis_agent acronis_scheduler; do
        service "$svc" stop 2>/dev/null || true
    done
}

unload_snapapi_modules() {
    log_info "Unloading SnapAPI kernel modules..."
    
    rmmod snapapi26 2>/dev/null || true
    rmmod snumbd26 2>/dev/null || true
    
    # Verify modules are unloaded
    if lsmod | grep -q snapapi26; then
        log_warn "snapapi26 module is still loaded (may be in use)"
    fi
}

remove_old_snapapi() {
    if [ -n "$CURRENT_VERSION" ]; then
        log_info "Removing old SnapAPI version $CURRENT_VERSION from DKMS..."
        dkms remove -m snapapi26 -v "$CURRENT_VERSION" --all 2>/dev/null || true
    fi
    
    # Also try to remove any other versions
    for ver in $(dkms status snapapi26 2>/dev/null | sed -n 's/snapapi26[,/]\s*\([0-9.]*\).*/\1/p'); do
        log_info "Removing SnapAPI version $ver from DKMS..."
        dkms remove -m snapapi26 -v "$ver" --all 2>/dev/null || true
    done
}

install_rpm_package() {
    local rpm_file="$1"
    
    log_info "Installing SnapAPI RPM package..."
    
    if [ "$USE_ACRONIS_RPMDB" = true ]; then
        # dpkg exists - use Acronis embedded RPM database
        local acronis_rpmdb="/var/lib/Acronis/.rpmdb"
        
        if [ ! -d "$acronis_rpmdb" ]; then
            log_warn "Acronis RPM database not found at $acronis_rpmdb, creating..."
            mkdir -p "$acronis_rpmdb"
        fi
        
        log_info "Installing RPM using Acronis RPM database..."
        rpm -U --force --force-debian --nodeps --nopostun "$rpm_file" \
            --dbpath "$acronis_rpmdb"
    else
        # Native RPM system - use system RPM database
        log_info "Installing RPM using system RPM database..."
        rpm -U --force --nodeps --nopostun "$rpm_file"
    fi
}

build_snapapi_module() {
    local kernel_ver=$(uname -r)
    local tarball="/usr/lib/Acronis/kernel_modules/snapapi26-${NEW_VERSION}-all.tar.gz"
    
    if [ ! -f "$tarball" ]; then
        log_error "SnapAPI tarball not found at: $tarball"
        log_error "The RPM package may not have been installed correctly"
        exit 1
    fi
    
    log_info "Loading SnapAPI sources into DKMS..."
    dkms ldtarball --force --archive "$tarball"
    
    log_info "Building and installing SnapAPI module for kernel $kernel_ver..."
    if ! dkms install -m snapapi26 -v "$NEW_VERSION"; then
        log_error "DKMS install failed"
        log_error "Check /var/lib/dkms/snapapi26/${NEW_VERSION}/build/make.log for details"
        exit 1
    fi
}

load_snapapi_modules() {
    log_info "Loading SnapAPI kernel modules..."
    
    modprobe snapapi26 || {
        log_error "Failed to load snapapi26 module"
        exit 1
    }
    
    # snumbd26 is optional, don't fail if it doesn't load
    modprobe snumbd26 2>/dev/null || true
}

start_acronis_services() {
    log_info "Starting Acronis services..."
    
    # Try systemd first
    for svc in acronis_mms acronis_agent acronis_scheduler; do
        systemctl start "$svc" 2>/dev/null || true
    done
    
    # Fallback to service command
    for svc in acronis_mms acronis_agent acronis_scheduler; do
        service "$svc" start 2>/dev/null || true
    done
}

verify_installation() {
    log_info "Verifying installation..."
    
    echo ""
    echo "=== DKMS Status ==="
    dkms status snapapi26
    
    echo ""
    echo "=== Loaded Modules ==="
    lsmod | grep -E "(snapapi|snumbd)" || log_warn "No SnapAPI modules loaded"
    
    echo ""
    echo "=== Module Info ==="
    modinfo snapapi26 2>/dev/null | grep -E "(filename|version|description)" || true
    
    echo ""
    if command -v acrocmd &>/dev/null; then
        echo "=== Disk List Test ==="
        acrocmd list disks 2>/dev/null || log_warn "acrocmd list disks failed"
    fi
}

# ============================================================================
# Mode-specific functions
# ============================================================================

replace_source_file() {
    local src_file="$1"
    local dkms_src_dir="${DKMS_SOURCE_DIR}-${CURRENT_VERSION}"
    local dkms_src="${dkms_src_dir}/snapapi26.c"
    
    if [ ! -d "$dkms_src_dir" ]; then
        log_error "DKMS source directory not found: $dkms_src_dir"
        log_error "Make sure SnapAPI is installed via DKMS first"
        exit 1
    fi
    
    if [ ! -f "$dkms_src" ]; then
        log_error "Source file not found: $dkms_src"
        exit 1
    fi
    
    log_info "Backing up original snapapi26.c..."
    cp "$dkms_src" "${dkms_src}.backup"
    
    log_info "Replacing snapapi26.c with $src_file..."
    cp "$src_file" "$dkms_src"
}

rebuild_current() {
    local kernel_ver=$(uname -r)
    
    log_info "Removing current build for kernel $kernel_ver..."
    dkms remove -m snapapi26 -v "$CURRENT_VERSION" -k "$kernel_ver" 2>/dev/null || true
    
    log_info "Building and installing SnapAPI module for kernel $kernel_ver..."
    if ! dkms install -m snapapi26 -v "$CURRENT_VERSION"; then
        log_error "DKMS install failed"
        log_error "Check /var/lib/dkms/snapapi26/${CURRENT_VERSION}/build/make.log for details"
        exit 1
    fi
}

# ============================================================================
# Main
# ============================================================================

# Parse arguments
case "${1:-}" in
    --rpm)
        MODE="rpm"
        RPM_FILE="${2:-}"
        if [ -z "$RPM_FILE" ]; then
            log_error "--rpm requires a file argument"
            usage
        fi
        if [ ! -f "$RPM_FILE" ]; then
            log_error "RPM file not found: $RPM_FILE"
            exit 1
        fi
        ;;
    --file)
        MODE="file"
        SOURCE_FILE="${2:-}"
        if [ -z "$SOURCE_FILE" ]; then
            log_error "--file requires a file argument"
            usage
        fi
        if [ ! -f "$SOURCE_FILE" ]; then
            log_error "Source file not found: $SOURCE_FILE"
            exit 1
        fi
        ;;
    --rebuild)
        MODE="rebuild"
        ;;
    -h|--help)
        usage
        ;;
    *)
        usage
        ;;
esac

log_info "=========================================="
log_info "SnapAPI Update Script"
log_info "=========================================="
log_info "Mode: $MODE"
log_info "Kernel: $(uname -r)"
log_info "=========================================="

check_root
detect_system
check_prerequisites
get_current_snapapi_version

# Mode-specific validation
case "$MODE" in
    rpm)
        extract_rpm_version "$RPM_FILE"
        log_info "RPM file: $RPM_FILE"
        log_info "New version: $NEW_VERSION"
        ;;
    file|rebuild)
        if [ -z "$CURRENT_VERSION" ]; then
            log_error "No SnapAPI version installed via DKMS"
            log_error "Use --rpm to install first"
            exit 1
        fi
        NEW_VERSION="$CURRENT_VERSION"
        if [ "$MODE" = "file" ]; then
            log_info "Source file: $SOURCE_FILE"
        fi
        ;;
esac

echo ""
case "$MODE" in
    rpm)
        log_warn "This will:"
        log_warn "  1. Stop Acronis services"
        log_warn "  2. Unload current SnapAPI modules"
        log_warn "  3. Remove old SnapAPI from DKMS"
        log_warn "  4. Install new SnapAPI RPM"
        log_warn "  5. Build and install new SnapAPI module"
        log_warn "  6. Load new modules and start services"
        ;;
    file)
        log_warn "This will:"
        log_warn "  1. Stop Acronis services"
        log_warn "  2. Unload current SnapAPI modules"
        log_warn "  3. Replace snapapi26.c in DKMS source"
        log_warn "  4. Rebuild SnapAPI module"
        log_warn "  5. Load new modules and start services"
        ;;
    rebuild)
        log_warn "This will:"
        log_warn "  1. Stop Acronis services"
        log_warn "  2. Unload current SnapAPI modules"
        log_warn "  3. Rebuild SnapAPI module"
        log_warn "  4. Load new modules and start services"
        ;;
esac

echo ""
read -p "Continue? [y/N] " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    log_info "Aborted by user"
    exit 0
fi

# Execute based on mode
stop_acronis_services
unload_snapapi_modules

case "$MODE" in
    rpm)
        remove_old_snapapi
        install_rpm_package "$RPM_FILE"
        build_snapapi_module
        ;;
    file)
        replace_source_file "$SOURCE_FILE"
        rebuild_current
        ;;
    rebuild)
        rebuild_current
        ;;
esac

load_snapapi_modules
start_acronis_services

echo ""
log_info "=========================================="
log_info "SnapAPI update completed successfully!"
log_info "=========================================="

verify_installation

echo ""
log_info "If you encounter issues, check:"
log_info "  - /var/lib/dkms/snapapi26/${NEW_VERSION}/build/make.log"
log_info "  - dmesg | grep -i snapapi"
log_info "  - journalctl -u acronis_mms"