#!/bin/bash

# Enable strict error handling
set -E -o pipefail
# ===================== Fixed Version Combinations =====================
declare -A VERSION_MAP=(
    # CUDA 12.1.1 compatible versions
    ["12.1.1_DRIVER"]="535.154.05"
    ["12.1.1_CUDNN"]="8.9.7.29"
    ["12.1.1_CUDNN_TAR"]="cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz"

    # CUDA 12.4.1 compatible versions
    ["12.4.1_DRIVER"]="550.90.07"
    ["12.4.1_CUDNN"]="9.2.0.82"
    ["12.4.1_CUDNN_TAR"]="cudnn-linux-x86_64-9.2.0.82_cuda12-archive.tar.xz"

    # CUDA 12.8.1 compatible versions
    ["12.8.1_DRIVER"]="570.133.20"
    ["12.8.1_CUDNN"]="9.8.0.87"
    ["12.8.1_CUDNN_TAR"]="cudnn-linux-x86_64-9.8.0.87_cuda12-archive.tar.xz"
)


get_cuda_filename(){
    if [ "$CUDA_VERSION" == "12.1.1" ]; then
        cuda_file="cuda_12.1.1_530.30.02_linux.run"
    elif [ "$CUDA_VERSION" == "12.4.1" ]; then
        cuda_file="cuda_12.4.1_550.54.15_linux.run"
    elif [ "$CUDA_VERSION" == "12.8.1" ]; then
        cuda_file="cuda_12.8.1_570.124.06_linux.run"
    else
        log_err "Unsupported CUDA version: $CUDA_VERSION"
        return 1
    fi
}

yum_source="http://internal.mirrors.ctyun.cn/ctyun-tools/ecs/gpu/tools/x86/nvidia/"
yum_source_back="http://169.254.169.253/ctyun-tools/ecs/gpu/tools/x86/nvidia/"

# Installation directory configuration
AUTO_INSTALL="/root/nvidia_auto_install"
DRIVER_DIR="${AUTO_INSTALL}/driver"
CUDA_DIR="${AUTO_INSTALL}/cuda"
CUDNN_DIR="${AUTO_INSTALL}/cudnn"
DEP_DIR="${AUTO_INSTALL}/dep"

# Process and file name configuration
DRIVER_PROCESS_NAME="${DRIVER_DIR}/NVIDIA-Linux-x86_64"
CUDA_PROCESS_NAME="${CUDA_DIR}/./cuda"
CUDNN_PROCESS_NAME="${CUDNN_DIR}/cudnn"
DOWNLOAD_DRIVER="NVIDIA-Linux"
DOWNLOAD_PROCESS_NAME="curl"

PROCESS_NAME="$0"
LOG_LEVEL=1
LOG_FILE="${AUTO_INSTALL}/debug.log"

# Command line argument variables
DRIVER_OVERRIDE=""
CHECK_MODE=false

# Create log directory immediately to avoid log function write failures
mkdir -p "${AUTO_INSTALL}" 2>/dev/null


init_dir(){
    mkdir -p "${AUTO_INSTALL}"
    mkdir -p "${DRIVER_DIR}"
    mkdir -p "${CUDA_DIR}"
    mkdir -p "${CUDNN_DIR}"
    mkdir -p "${DEP_DIR}"
}
log_debug() {
  content="[DEBUG] $(date '+%Y-%m-%d %H:%M:%S') $@"
  [ $LOG_LEVEL -le 1  ] && echo -e "\033[32m"  ${content}  "\033[0m"
  [ $LOG_LEVEL -le 1  ] && echo -e "\033[32m"  ${content}  "\033[0m" >> ${LOG_FILE} 2>&1
}

log_info() {
  content="[INFO] $(date '+%Y-%m-%d %H:%M:%S') $@"
  [ $LOG_LEVEL -le 2  ] && echo -e "\033[32m"  ${content} "\033[0m"
  [ $LOG_LEVEL -le 2  ] && echo -e "\033[32m"  ${content} "\033[0m" >> ${LOG_FILE} 2>&1
}

log_warn() {
  content="[WARN] $(date '+%Y-%m-%d %H:%M:%S') $@"
  [ $LOG_LEVEL -le 3  ] && echo -e "\033[33m" ${content} "\033[0m"
  [ $LOG_LEVEL -le 3  ] && echo -e "\033[33m" ${content} "\033[0m" >> ${LOG_FILE} 2>&1
}

log_err() {
  content="[ERROR] $(date '+%Y-%m-%d %H:%M:%S') $@"
  [ $LOG_LEVEL -le 4  ] && echo -e "\033[31m" ${content} "\033[0m"
  [ $LOG_LEVEL -le 4  ] && echo -e "\033[31m" ${content} "\033[0m" >> ${LOG_FILE} 2>&1
}

log_always() {
   content="[ALWAYS] $(date '+%Y-%m-%d %H:%M:%S') $@"
   [ $LOG_LEVEL -le 5  ] && echo -e  "\033[32m" ${content} "\033[0m"
   [ $LOG_LEVEL -le 5  ] && echo -e  "\033[32m" ${content} "\033[0m" >> ${LOG_FILE} 2>&1
}

log_only() {
  local content="[ONLY] $(date '+%Y-%m-%d %H:%M:%S') $*"
  [ $LOG_LEVEL -le 3 ] && echo -e "${content}" >> "${LOG_FILE}" 2>&1
}

# Parse command line arguments
parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --driver)
                if [ -z "$2" ] || [[ "$2" == --* ]]; then
                    log_err "--driver requires a version argument"
                    log_err "Usage: $0 [--driver <version>] [check]"
                    log_err "Example: $0 --driver 550.90.07"
                    exit 1
                fi
                DRIVER_OVERRIDE="$2"
                shift 2
                ;;
            check)
                CHECK_MODE=true
                shift
                ;;
            --help|-h|-help)
                echo ""
                echo "Usage: $0 [OPTIONS] [check]"
                echo ""
                echo "Options:"
                echo "  --driver <version>  Specify driver version"
                echo "  check                Check installation status"
                echo "  --help               Show this help message"
                echo ""
                echo "Examples:"
                echo "  $0                               Auto install with recommended version"
                echo "  $0 --driver 550.90.07         Install with specific driver version"
                echo "  $0 check                        Check GPU installation status"
                echo ""
                echo "Supported driver versions:"
                echo "  535.154.05  (CUDA 12.1.1, cuDNN 8.9.7.29)"
                echo "  550.90.07   (CUDA 12.4.1, cuDNN 9.2.0.82)"
                echo "  570.133.20   (CUDA 12.8.1, cuDNN 9.8.0.87)"
                echo ""
                exit 0
                ;;
            *)
                log_err "Unknown argument: $1"
                log_err "Usage: $0 [--driver <version>] [check] [--help]"
                exit 1
                ;;
        esac
    done
}

# Version validation functions
get_cuda_by_driver() {
    case "$1" in
        "535.154.05") echo "12.1.1" ;;
        "550.90.07")  echo "12.4.1" ;;
        "570.133.20") echo "12.8.1" ;;
        *) echo "" ;;
    esac
}

validate_driver_version() {
    local driver="$1"
    local matched_cuda=""
    
    matched_cuda=$(get_cuda_by_driver "$driver")
    
    if [ -z "$matched_cuda" ]; then
        log_err "Unsupported driver version: $driver"
        log_err ""
        log_err "Supported version combinations:"
        log_err "  Driver        | CUDA    | cuDNN"
        log_err "  --------------|---------|--------"
        log_err "  535.154.05    | 12.1.1  | 8.9.7.29"
        log_err "  550.90.07     | 12.4.1  | 9.2.0.82"
        log_err "  570.133.20    | 12.8.1  | 9.8.0.87"
        log_err ""
        log_err "Please specify a valid driver version from the list above."
        log_err "Example: bash $0 --driver 550.90.07"
        return 1
    fi
    return 0
}

# Version selection logic
select_version() {
    # Detect operating system information
    if [ -f "/etc/os-release" ]; then
        os=$(. /etc/os-release; echo $ID)
        os_version=$(. /etc/os-release; echo $VERSION_ID)
    else
        log_warn "/etc/os-release does not exist, trying to detect from /etc/issue"
        if grep -q "CentOS" /etc/issue; then
            os="centos"
            os_version=$(grep -oE "[0-9]+\.[0-9]+" /etc/issue | head -1)
        else
            log_err "Unable to detect operating system type"
            return 1
        fi
    fi
    kernel_version=$(uname -r | cut -d'-' -f1)
    major_kernel=$(echo $kernel_version | cut -d'.' -f1-2)

    log_info "Operating System: $os $os_version"
    log_info "Kernel Version: $kernel_version"

    # Fine-tune recommended version based on OS
    if [[ "$os" == "ubuntu" ]]; then
        log_info "OS version.kernel_major: $os_version.$major_kernel"
        case "$os_version.$major_kernel" in
            "18.04.4.15")
                RECOMMENDED_CUDA="12.4.1"
                SUPPORTED_CUDA=("12.1.1" "12.4.1")
                ;;
            "20.04.5.4")
                RECOMMENDED_CUDA="12.8.1"
                SUPPORTED_CUDA=("12.1.1" "12.4.1" "12.8.1")
                ;;
            "22.04.5.15")
                RECOMMENDED_CUDA="12.8.1"
                SUPPORTED_CUDA=("12.4.1" "12.8.1")
                ;;
            "24.04.6.8")
                RECOMMENDED_CUDA="12.8.1"
                SUPPORTED_CUDA=("12.1.1" "12.4.1" "12.8.1")
                ;;
            *)
                log_err "Unsupported Ubuntu version or kernel: $os_version $kernel_version"
                exit 1
                ;;
        esac
    elif [[ "$os" == "ctyunos" ]]; then
        log_info "OS version.kernel_major: $os_version.$major_kernel"
        case "$os_version.$major_kernel" in
            "2.0.1.4.19")
                RECOMMENDED_CUDA="12.8.1"
                SUPPORTED_CUDA=("12.1.1" "12.4.1" "12.8.1")
                ;;
            "22.06.4.19")
                RECOMMENDED_CUDA="12.8.1"
                SUPPORTED_CUDA=("12.1.1" "12.4.1" "12.8.1")
                ;;
            "23.01.5.10")
                RECOMMENDED_CUDA="12.8.1"
                SUPPORTED_CUDA=("12.1.1" "12.4.1" "12.8.1")
                ;;
            *)
                log_err "Unsupported CTyunOS version or kernel: $os_version $kernel_version"
                exit 1
                ;;
        esac

    elif [[ "$os" == "centos" ]]; then
        log_info "OS version.kernel_major: $os_version.$major_kernel"
        case "$os_version.$major_kernel" in
            "7.3.10")
                RECOMMENDED_CUDA="12.4.1"
                SUPPORTED_CUDA=("12.1.1" "12.4.1")
                ;;
            "8.4.18")
                RECOMMENDED_CUDA="12.8.1"
                SUPPORTED_CUDA=("12.1.1" "12.4.1" "12.8.1")
                ;;
            "9.5.14")
                RECOMMENDED_CUDA="12.8.1"
                SUPPORTED_CUDA=("12.8.1")
                ;;
            *)
                log_err "Unsupported CentOS version or kernel: $os_version $kernel_version"
                exit 1
                ;;
        esac
    elif [[ "$os" == "kylin" ]]; then
        log_info "OS version.kernel_major: $os_version.$major_kernel"
        case "$os_version.$major_kernel" in
            "V10.4.19")
                RECOMMENDED_CUDA="12.8.1"
                SUPPORTED_CUDA=("12.1.1" "12.4.1" "12.8.1")
                ;;
            *)
                log_err "Unsupported Kylin version or kernel: $os_version $kernel_version"
                exit 1
                ;;
        esac
    else
        log_err "Unsupported operating system, please refer to the documentation"
        exit 1
    fi


    CUDA_VERSION=$RECOMMENDED_CUDA

    # If user specified driver version, override automatic selection
    if [ -n "$DRIVER_OVERRIDE" ]; then
        if validate_driver_version "$DRIVER_OVERRIDE"; then
            DRIVER_VERSION="$DRIVER_OVERRIDE"
            CUDA_VERSION=$(get_cuda_by_driver "$DRIVER_OVERRIDE")
            
            # Check if the CUDA version is supported by the current OS
            local cuda_supported=false
            for supported_cuda in "${SUPPORTED_CUDA[@]}"; do
                if [ "$CUDA_VERSION" = "$supported_cuda" ]; then
                    cuda_supported=true
                    break
                fi
            done
            
            if [ "$cuda_supported" = false ]; then
                log_err ""
                log_err "==========================================================================="
                log_err "ERROR: Driver $DRIVER_VERSION (CUDA $CUDA_VERSION) is NOT supported on this OS!"
                log_err "==========================================================================="
                log_err ""
                log_err "Current OS: $os $os_version (kernel $kernel_version)"
                log_err "Requested driver: $DRIVER_VERSION -> CUDA $CUDA_VERSION"

                log_err "Recommended driver versions for this OS:"
                local recommended_driver=${VERSION_MAP["${RECOMMENDED_CUDA}_DRIVER"]}
                log_err "  - Driver $recommended_driver (CUDA $RECOMMENDED_CUDA) [RECOMMENDED]"
                for supported_cuda in "${SUPPORTED_CUDA[@]}"; do
                    local driver_for_cuda=${VERSION_MAP["${supported_cuda}_DRIVER"]}
                    if [ "$driver_for_cuda" != "$recommended_driver" ]; then
                        log_err "  - Driver $driver_for_cuda (CUDA $supported_cuda)"
                    fi
                done
                log_err ""
                log_err "Please use one of the supported drivers listed above."
                log_err "Example: bash $0 --driver $recommended_driver"
                log_err ""
                exit 1
            fi
            
            CUDA_SHORT_VERSION=$(echo $CUDA_VERSION | cut -d'.' -f1-2)
            CUDNN_VERSION=${VERSION_MAP["${CUDA_VERSION}_CUDNN"]}
            CUDNN_TAR=${VERSION_MAP["${CUDA_VERSION}_CUDNN_TAR"]}
            log_info "User specified driver version: $DRIVER_VERSION"
            log_info "Matched CUDA version: $CUDA_VERSION"
            log_info "Matched cuDNN version: $CUDNN_VERSION"
        else
            exit 1
        fi
    else
        CUDA_SHORT_VERSION=$(echo $CUDA_VERSION | cut -d'.' -f1-2)
        DRIVER_VERSION=${VERSION_MAP["${CUDA_VERSION}_DRIVER"]}
        CUDNN_VERSION=${VERSION_MAP["${CUDA_VERSION}_CUDNN"]}
        CUDNN_TAR=${VERSION_MAP["${CUDA_VERSION}_CUDNN_TAR"]}
    fi

    log_info "=== Selected Version ==="
    log_info "Operating System: $os $os_version "
    log_info "Kernel Version: $kernel_version"
    log_info "CUDA Version: $CUDA_VERSION"
    log_info "Driver Version: $DRIVER_VERSION"
    log_info "cuDNN Version: $CUDNN_VERSION"
}


check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_err "This script must be run with root privileges, please use sudo or switch to root user"
        exit 1
    fi
}


check_and_install_commands() {
    local REQUIRED_COMMANDS=("$@")
    local MISSING_COMMANDS=()
    local PACKAGE_MANAGER=""
    local INSTALL_CMD=""

    if command -v yum &>/dev/null; then
        PACKAGE_MANAGER="yum"
        INSTALL_CMD="yum install -y"
    elif command -v apt-get &>/dev/null; then
        PACKAGE_MANAGER="apt-get"
        INSTALL_CMD="apt-get install -y"
    elif command -v dnf &>/dev/null; then
        PACKAGE_MANAGER="dnf"
        INSTALL_CMD="dnf install -y"
    else
        log_err "No supported package manager found (yum, apt-get, dnf)"
        return 1
    fi

    log_debug "Using package manager: $PACKAGE_MANAGER"

    for command in "${REQUIRED_COMMANDS[@]}"; do
        if ! command -v "$command" &>/dev/null; then
            log_warn "Command $command not found, will try to install"
            MISSING_COMMANDS+=("$command")
        else
            log_debug "Command $command is already installed"
        fi
    done

    if [ ${#MISSING_COMMANDS[@]} -eq 0 ]; then
        log_info "All required commands are installed"
        return 0
    fi

    log_info "Missing commands: ${MISSING_COMMANDS[*]}"

    for command in "${MISSING_COMMANDS[@]}"; do
        log_info "Installing $command..."

        case "$command" in
            "gcc")
                package_name="gcc"
                ;;
            "make")
                package_name="make"
                ;;
            "wget")
                package_name="wget"
                ;;
            "tar")
                package_name="tar"
                ;;
            "curl")
                package_name="curl"
                ;;
            "bzip2")
                package_name="bzip2"
                ;;
            *)
                package_name="$command"
                log_warn "Package name for command $command is unknown, using $command as package name"
                ;;
        esac

        if ! $INSTALL_CMD "$package_name"; then
            log_err "Failed to install $command (package name: $package_name)"
            return 1
        fi

        log_info "$command installed successfully"
    done

    for command in "${MISSING_COMMANDS[@]}"; do
        if ! command -v "$command" &>/dev/null; then
            log_err "Command $command still not available after installation"
            return 1
        fi
    done

    log_info "All required commands have been successfully installed"
    return 0
}

install_kernel_redhat_series() {
    kernel_version=$(uname -r)
    log_info "Current kernel version: $kernel_version"

    # Check if kernel-devel is already installed
    if ! rpm -q kernel-devel-$kernel_version &>/dev/null; then
        log_info "Installing kernel-devel-$kernel_version ..."

        # Check if yum is available
        if ! command -v yum &>/dev/null; then
            log_err "YUM package manager not available, cannot install kernel-devel"
            return 1
        fi

        # Install kernel-devel
        if ! yum install -y kernel-devel-$kernel_version; then
            log_err "Failed to install kernel-devel-$kernel_version! Please check network connection and package manager status"
            return 1
        fi
        log_info "kernel-devel-$kernel_version installed successfully"
    else
        log_info "kernel-devel-$kernel_version is already installed"
    fi

    # Check OS version
    if [ -f "/etc/os-release" ]; then
        os_version=$(. /etc/os-release; echo $VERSION_ID)
        if [[ "$os_version" == "8"* ]]; then
            log_info "Operating system is CentOS 8 or similar version"

            # Check if elfutils packages are already installed
            if ! rpm -q elfutils-libelf-devel-static &>/dev/null && ! rpm -q elfutils-libelf-devel &>/dev/null; then
                log_info "Installing required elfutils packages ..."

                # Try installing first, if failed retry with --allowerasing
                if ! yum install -y elfutils-libelf-devel elfutils-libelf-devel-static; then
                    log_warn "First installation attempt failed, trying forced replacement of conflicting packages..."
                    if ! yum install -y elfutils-libelf-devel elfutils-libelf-devel-static --allowerasing; then
                        log_err "Failed to install elfutils packages! Please check network connection and package manager status"
                        return 1
                    fi
                fi
                log_info "elfutils packages installed successfully"
            else
                log_info "Required elfutils packages are already installed"
            fi
        fi
    fi
    return 0
}

install_kernel_debian_series() {
    kernel_version=$(uname -r)
    log_info "Current kernel version: $kernel_version"

    # Check if linux-headers is already installed
    if ! dpkg -l | grep -q "linux-headers-${kernel_version}"; then
        log_info "Installing linux-headers-${kernel_version} ..."

        # Check if apt-get is available
        if ! command -v apt-get &>/dev/null; then
            log_err "APT package manager not available, cannot install linux-headers"
            return 1
        fi

        # Ensure package list is up to date
        log_info "Updating package list..."
        if ! apt-get update > ${AUTO_INSTALL}/apt_update.log 2>&1; then
            log_warn "Failed to update package list, but will try to continue installing linux-headers"
        fi

        # Install linux-headers
        if ! apt-get install -y --allow-unauthenticated linux-headers-${kernel_version} > ${AUTO_INSTALL}/kernel_install.log 2>&1; then
            log_err "Failed to install linux-headers-${kernel_version}! Please check ${AUTO_INSTALL}/kernel_install.log for details"
            return 1
        fi
        log_info "linux-headers-${kernel_version} installed successfully"
    else
        log_info "linux-headers-${kernel_version} is already installed"
    fi
    return 0
}

nouveau_disable_redhat_series() {
    log_info "Starting to disable nouveau driver for RedHat series systems..."

    # Create blacklist file
    if  [ ! -f /etc/modprobe.d/blacklist-nouveau.conf ];then
        log_info "Creating /etc/modprobe.d/blacklist-nouveau.conf file"
        if ! echo "blacklist nouveau" > /etc/modprobe.d/blacklist-nouveau.conf; then
            log_err "Failed to create blacklist file, insufficient permissions"
            return 1
        fi
        if ! echo "options nouveau modeset=0" >> /etc/modprobe.d/blacklist-nouveau.conf; then
            log_err "Failed to add content to blacklist file"
            return 1
        fi
        log_info "Created /etc/modprobe.d/blacklist-nouveau.conf to disable nouveau driver"
    else
        log_info "/etc/modprobe.d/blacklist-nouveau.conf file already exists"
    fi

    # Check and remove nouveau module
    if lsmod | grep -q nouveau; then
        log_info "Removing nouveau module..."
        rmmod nouveau || {
            log_err "Failed to remove nouveau module, may need to reboot the system"
            return 1
        }

        # Recreate initramfs
        log_info "Executing dracut --force to recreate kernel initial RAM filesystem..."
        if ! dracut --force; then
            log_err "Failed to recreate initial RAM filesystem"
            return 1
        fi

        # Verify nouveau is unloaded
        if lsmod | grep -q nouveau; then
            log_err "Nouveau driver is still loaded, system reboot required"
            return 1
        fi
        log_info "nouveau module successfully removed"
    else
        log_info "Nouveau driver not loaded, no action needed"
    fi

    log_info "Nouveau driver successfully disabled for RedHat series systems"
    return 0
}

nouveau_disable_debian_series() {
    log_info "Starting to disable nouveau driver for Debian series systems..."

    # Create blacklist file
    if  [ ! -f /etc/modprobe.d/blacklist-nouveau.conf ];then
        log_info "Creating /etc/modprobe.d/blacklist-nouveau.conf file"
        if ! echo "blacklist nouveau" > /etc/modprobe.d/blacklist-nouveau.conf; then
            log_err "Failed to create blacklist file, insufficient permissions"
            return 1
        fi
        if ! echo "blacklist lbm-nouveau" >> /etc/modprobe.d/blacklist-nouveau.conf; then
            log_err "Failed to add content to blacklist file"
            return 1
        fi
        if ! echo "options nouveau modeset=0" >> /etc/modprobe.d/blacklist-nouveau.conf; then
            log_err "Failed to add content to blacklist file"
            return 1
        fi
        log_info "Created /etc/modprobe.d/blacklist-nouveau.conf to disable nouveau driver"
    else
        log_info "/etc/modprobe.d/blacklist-nouveau.conf file already exists"
    fi

    # Check and remove nouveau module
    if lsmod | grep -q nouveau; then
        log_info "Removing nouveau module..."
        rmmod nouveau || {
            log_err "Failed to remove nouveau module, may need to reboot the system"
            return 1
        }

        # Recreate initramfs
        log_info "Executing update-initramfs -u to recreate kernel initial RAM filesystem..."
        if ! update-initramfs -u; then
            log_err "Failed to recreate initial RAM filesystem"
            return 1
        fi

        # Verify nouveau is unloaded
        if lsmod | grep -q nouveau; then
            log_err "Nouveau driver is still loaded, system reboot required"
            return 1
        fi
        log_info "nouveau module successfully removed"
    else
        log_info "Nouveau driver not loaded, no action needed"
    fi

    log_info "Nouveau driver successfully disabled for Debian series systems"
    return 0
}

# Disable nouveau driver
disable_nouveau() {
    if [ -z "$os" ]; then
        log_err "Install error: os variable is undefined!"
        exit 1
    fi

    case "$os" in
        "ubuntu")
            # Ubuntu-specific logic: disable nouveau + version detection + install debian series kernel
            nouveau_disable_debian_series
            if [ $? -ne 0 ]; then
                log_err "Install error: Failed to disable nouveau for Debian series!"
                exit 1
            fi

            # Get Ubuntu version and convert to shorthand (e.g. 18.04 -> ubuntu1804)
            version=$(. /etc/os-release; echo "$VERSION_ID")
            case "$version" in
                "16.04") ubuntu_version="ubuntu1604" ;;
                "18.04") ubuntu_version="ubuntu1804" ;;
                "20.04") ubuntu_version="ubuntu2004" ;;
                "22.04") ubuntu_version="ubuntu2204" ;;
                "24.04") ubuntu_version="ubuntu2404" ;;
                *)
                    log_err "Error: Unsupported Ubuntu version $version!"
                    exit 1
                    ;;
            esac

            install_kernel_debian_series
            if [ $? -ne 0 ]; then
                log_err "Install error: kernel-devel installation failed!"
                exit 1
            fi
            ;;

        "centos"|"ctyunos"|"kylin")
            # Merge duplicate logic for CentOS/CTyunOS/Kylin
            nouveau_disable_redhat_series
            if [ $? -ne 0 ]; then
                log_err "Install error: Failed to disable nouveau for RedHat series!"
                exit 1
            fi

            # Keep CentOS original VERSION_ID read logic (even if unused, for compatibility)
            if [ "$os" = "centos" ] && [ -f "/etc/os-release" ]; then
                version=$(. /etc/os-release; echo "$VERSION_ID")
            fi

            install_kernel_redhat_series
            if [ $? -ne 0 ]; then
                log_err "Install error: Failed to install kernel development package for RedHat series!"
                exit 1
            fi
            ;;

        *)
            log_err "Install error: Invalid operating system! Installation failed! OS is $os"
            exit 1
            ;;
    esac
}

download_file() {
    url=$1
    expected_size=$2
    expected_md5=$3
    num_of_retries=12

    local filename=$(basename "$url")
    local file_path="${PWD}/${filename}"
    local valid_expected_size=0

    if [ -f "$file_path" ]; then
        expected_size=${expected_size:-$(get_expected_file_size "$url")}
        valid_expected_size=$([ "$expected_size" -gt 0 ] && echo 1 || echo 0)

        if validate_file_integrity "$file_path" "$expected_size" "$expected_md5"; then
            log_info "File already exists and integrity verified, no need to download: $filename"
            return 0
        else
            log_warn "File exists but integrity verification failed, deleting and re-downloading: $filename"
            rm -f "$file_path"
        fi
    fi

    if [ -z "$expected_size" ] || [ "$expected_size" -le 0 ]; then
        expected_size=$(get_expected_file_size "$url")
        valid_expected_size=$([ "$expected_size" -gt 0 ] && echo 1 || echo 0)
    else
        valid_expected_size=1
        log_info "Using provided expected file size: $expected_size bytes"
    fi

    local downloaded_size=0
    if [ -f "$file_path" ]; then
        downloaded_size=$(stat -c%s "$file_path" 2>/dev/null || stat -f%z "$file_path" 2>/dev/null)
        log_info "Detected partially downloaded file, size: ${downloaded_size} bytes, will attempt resume"

        if [ "$valid_expected_size" -eq 1 ] && [ "$downloaded_size" -ge "$expected_size" ]; then
            log_warn "Downloaded file size exceeds expected, deleting and re-downloading"
            rm -f "$file_path"
            downloaded_size=0
        fi
    fi

    for (( i=1; i<=$num_of_retries; i++ )); do
        log_info "Starting download (attempt $i/$num_of_retries): $url"

        curl -L -C - -# -o "$file_path" "$url"
        local curl_exit_code=$?

        log_info "Download request completed, curl exit code: $curl_exit_code"

        if [ $curl_exit_code -eq 0 ]; then
            log_info "Download completed: $url (curl exit code: $curl_exit_code)"
            if validate_file_integrity "$file_path" "$expected_size" "$expected_md5"; then
                log_info "File integrity verified: $filename"
                return 0
            else
                log_err "File integrity verification failed: $filename"
                rm -f "$file_path"
                if [ $i -lt $num_of_retries ]; then
                    log_warn "Re-downloading corrupted file..."
                    sleep 5
                fi
            fi
        else

            local http_code=$(curl -L -s -o /dev/null -w "%{http_code}" --max-time 10 "$url")
            log_err "curl execution failed (exit code: $curl_exit_code), URL status code: $http_code"

            case "$http_code" in
                404) log_err "File not found (HTTP 404): $url" ;;
                403) log_err "Access denied (HTTP 403): $url" ;;
                416)
                    log_err "Resume failed (HTTP 416): $url, verifying file integrity..."
                    if validate_file_integrity "$file_path" "$expected_size" "$expected_md5"; then
                        log_info "File integrity verified, 416 is normal resume completion status"
                        return 0
                    elif [ "$valid_expected_size" -eq 1 ] && [ $(stat -c%s "$file_path") -eq "$expected_size" ]; then
                        log_info "File size matches expected, skipping deletion"
                    else
                        rm -f "$file_path"
                    fi
                    ;;
                000|"") log_err "Network connection failed/curl execution error (exit code: $curl_exit_code): $url" ;;
                *) log_err "HTTP error (status code: $http_code): $url" ;;
            esac

            if [ $i -lt $num_of_retries ]; then
                log_warn "Waiting 5 seconds before retry..."
                sleep 5
            fi
        fi
    done

    log_err "Download failed after $num_of_retries retries"
    return 1
}
get_expected_file_size() {
    local url="$1"
    local expected_size=0
    local curl_exit_code=0

    expected_size=$(curl -sIL --max-time 10 "$url" 2>/dev/null \
        | grep -i '^Content-Length:' \
        | tail -1 \
        | awk '{print $2}' \
        | tr -d '\r\n' \
        | sed 's/[^0-9]//g')  # Fix: complete sed replacement syntax

    curl_exit_code=${PIPESTATUS[0]}

    if [ $curl_exit_code -ne 0 ]; then
        log_only "Failed to get file size: curl execution error (URL: $url, exit code: $curl_exit_code)"
        echo 0
        return 1
    fi

    if [ -z "$expected_size" ] || [ "$expected_size" -le 0 ]; then
        log_only "Failed to get file size: URL has no valid Content-Length response header (URL: $url)"
        echo 0
        return 1
    fi

    log_only "Successfully retrieved file size: $expected_size bytes (URL: $url)"
    echo "$expected_size"
    return 0
}

download(){
    driver_file="NVIDIA-Linux-x86_64-${DRIVER_VERSION}.run"
    get_cuda_filename
    cudnn_file=${CUDNN_TAR}
    DRIVER_URL="${yum_source}${driver_file}"
    CUDA_URL="${yum_source}${cuda_file}"
    CUDNN_URL="${yum_source}${cudnn_file}"

    log_info "Starting to download driver file: ${driver_file}"
    cd ${DRIVER_DIR}
    download_file ${DRIVER_URL}
    if [ $? -ne 0 ]; then
        log_err "Download error: Driver download failed!"
        return 1
    fi

    log_info "Starting to download CUDA file: ${cuda_file}"
    cd ${CUDA_DIR}
    download_file ${CUDA_URL}
    if [ $? -ne 0 ]; then
        log_err "Download error: CUDA download failed!"
        return 1
    fi

    log_info "Starting to download cuDNN file: ${cudnn_file}"
    cd ${CUDNN_DIR}
    download_file ${CUDNN_URL}
    if [ $? -ne 0 ]; then
        log_err "Download error: cuDNN download failed!"
        return 1
    fi

    # Set executable permissions
    log_info "Setting file executable permissions..."
    chmod +x ${AUTO_INSTALL}/driver/* 2>/dev/null
    chmod +x ${AUTO_INSTALL}/cuda/* 2>/dev/null
    chmod +x ${AUTO_INSTALL}/cudnn/* 2>/dev/null

    log_info "All files downloaded and verified successfully"
    return 0
}

# Validate file integrity (supports file size and MD5 verification)
validate_file_integrity() {
    local file_path="$1"
    local expected_size="$2"
    local expected_md5="$3"

    # Check if file exists
    if [ ! -f "$file_path" ]; then
        log_err "$file_path file does not exist: $file_path"
        return 1
    fi

    # Check if file size is greater than 0
    local file_size=$(stat -c "%s" "$file_path" 2>/dev/null || stat -f "%z" "$file_path" 2>/dev/null)
    if [ -z "$file_size" ] || [ "$file_size" -eq 0 ]; then
        log_err "$file_path file is empty or unable to get size: $file_path"
        return 1
    fi

    # If expected size is provided, compare
    if [ -n "$expected_size" ] && [ "$expected_size" -gt 0 ]; then
        if [ "$file_size" -ne "$expected_size" ]; then
            log_err "$file_path file size mismatch: actual $file_size bytes, expected $expected_size bytes"
            return 1
        else
            log_info "$file_path file size verified: $file_size bytes"
        fi
    else
        log_info "$file_path file size verified: $file_size bytes"
    fi

    # If expected MD5 is provided, verify
    if [ -n "$expected_md5" ]; then
        # Check if MD5 tool is available
        if command -v md5sum &> /dev/null; then
            local actual_md5=$(md5sum "$file_path" | awk '{print $1}')
        elif command -v md5 &> /dev/null; then
            local actual_md5=$(md5 "$file_path" | awk '{print $4}')
        elif command -v openssl &> /dev/null; then
            local actual_md5=$(openssl md5 "$file_path" | awk '{print $2}')
        else
            log_warn "No MD5 tool available, skipping MD5 verification"
            return 0
        fi

        if [ "$actual_md5" = "$expected_md5" ]; then
            log_info "$file_path MD5 verified: $actual_md5"
        else
            log_err "$file_path MD5 verification failed: actual $actual_md5, expected $expected_md5"
            return 1
        fi
    else
        # Even without expected MD5, calculate and record actual MD5 value
        if command -v md5sum &> /dev/null; then
            local actual_md5=$(md5sum "$file_path" | awk '{print $1}')
            log_info "$file_path file MD5: $actual_md5"
        elif command -v md5 &> /dev/null; then
            local actual_md5=$(md5 "$file_path" | awk '{print $4}')
            log_info "$file_path file MD5: $actual_md5"
        elif command -v openssl &> /dev/null; then
            local actual_md5=$(openssl md5 "$file_path" | awk '{print $2}')
            log_info "$file_path file MD5: $actual_md5"
        else
            log_warn "$file_path No MD5 tool available, cannot calculate MD5 value"
        fi
        log_warn "$file_path No expected MD5 value, skipping MD5 verification"
    fi

    return 0
}

# Install driver
install_driver() {

    if command -v nvidia-smi &> /dev/null; then
        log_info "NVIDIA driver already installed, skipping installation"
        return 0
    fi

    # Check if driver file exists and is executable
    if [ ! -f "${AUTO_INSTALL}/driver/${driver_file}" ]; then
        log_err "NVIDIA driver installation error: Driver file ${AUTO_INSTALL}/driver/${driver_file} not found!"
        return 1
    fi

    if [ ! -x "${AUTO_INSTALL}/driver/${driver_file}" ]; then
        log_warn "Driver file is not executable, setting executable permission..."
        chmod +x "${AUTO_INSTALL}/driver/${driver_file}"
        if [ $? -ne 0 ]; then
            log_err "NVIDIA driver installation error: Failed to set driver file executable permission!"
            return 1
        fi
    fi

    # Execute driver installation
    log_info "Starting NVIDIA driver ${DRIVER_VERSION} installation..."
    ${AUTO_INSTALL}/driver/${driver_file} --silent  --dkms
    if [ $? -ne 0 ]; then
        log_err "NVIDIA driver installation error: Driver installation failed! Please check installation log for details"
        return 1
    fi

    # Verify driver installation success
    if command -v nvidia-smi &> /dev/null; then
        log_always "Driver $DRIVER_VERSION installed successfully!"
        log_info "NVIDIA driver installed successfully, version: $(nvidia-smi --query-gpu=driver_version --format=csv,noheader)"
        return 0
    else
        log_err "NVIDIA driver installation error: Driver installation completed, but nvidia-smi is not available!"
        return 1
    fi
}

# Install CUDA
install_cuda() {

    if command -v /usr/local/cuda/bin/nvcc &> /dev/null; then
        log_info "CUDA already installed, skipping installation"
        return 0
    fi

    # Check if CUDA directory exists
    if [ ! -d "${CUDA_DIR}" ]; then
        log_err "CUDA directory ${CUDA_DIR} does not exist"
        return 1
    fi

    cd ${CUDA_DIR} || {
        log_err "Failed to switch to directory ${CUDA_DIR}"
        return 1
    }

    # Find CUDA installation file
    cuda_file=$(find . -maxdepth 1 -name "cuda*${CUDA_VERSION}*" | head -1)
    log_info "CUDA file: $cuda_file"

    if [ -z "$cuda_file" ]; then
        log_err "CUDA installation error: CUDA file not found, installation failed!"
        return 1
    fi

    # Check if CUDA file is executable
    if [ ! -x "$cuda_file" ]; then
        log_warn "CUDA file is not executable, setting executable permission..."
        chmod +x "$cuda_file"
        if [ $? -ne 0 ]; then
            log_err "CUDA installation error: Failed to set CUDA file executable permission!"
            return 1
        fi
    fi

    # Execute CUDA installation
    log_info "Starting CUDA ${CUDA_VERSION} installation..."
    sh "$cuda_file" --silent --toolkit --samples --samplespath=/root
    if [ $? -ne 0 ]; then
        log_err "CUDA installation error: CUDA installation failed!"
        return 1
    fi

    # Install CUDA patches
    cuda_patchfiles=$(find . -maxdepth 1 -name "cuda*${CUDA_VERSION}*" ! -name "$(basename $cuda_file)")
    for cuda_patch in $cuda_patchfiles; do
        log_info "Installing CUDA patch file: $cuda_patch"
        # Check if patch file is executable
        if [ ! -x "$cuda_patch" ]; then
            log_warn "CUDA patch file is not executable, setting executable permission..."
            chmod +x "$cuda_patch"
            if [ $? -ne 0 ]; then
                log_err "CUDA installation error: Failed to set CUDA patch file executable permission!"
                return 1
            fi
        fi
        # Execute patch installation
        sh "$cuda_patch" --silent --installdir=/usr/local/cuda --accept-eula
        if [ $? -ne 0 ]; then
            log_err "CUDA installation error: CUDA patch installation failed!"
            return 1
        fi
    done

    # Verify CUDA installation success
    if command -v /usr/local/cuda/bin/nvcc &> /dev/null; then
        log_always "CUDA $CUDA_VERSION installed successfully!"
        log_info "CUDA version: $(/usr/local/cuda/bin/nvcc --version | grep release | awk '{print $NF}')"
        return 0
    else
        log_err "CUDA installation error: CUDA installation completed, but nvcc is not available!"
        return 1
    fi
}

# Install cuDNN
install_cudnn() {

    if [ -f /usr/local/cuda/include/cudnn.h ] || [ -f /usr/local/cuda/include/cudnn_version.h ]; then
        log_info "cuDNN already installed, skipping installation"
        return 0
    fi

    # Check if cuDNN file exists
    if [ ! -f "${AUTO_INSTALL}/cudnn/$cudnn_file" ]; then
        log_err "cuDNN installation error: cuDNN file ${AUTO_INSTALL}/cudnn/$cudnn_file not found!"
        return 1
    fi

    # Extract cuDNN file
    log_info "Starting to extract cuDNN file..."
    tar xvf ${AUTO_INSTALL}/cudnn/$cudnn_file -C ${AUTO_INSTALL}/cudnn
    if [ $? -ne 0 ]; then
        log_err "cuDNN installation error: Failed to extract cuDNN file!"
        return 1
    fi

    # Check if CUDA directory exists
    if [ ! -d "/usr/local/cuda/include" ] || [ ! -d "/usr/local/cuda/lib64" ]; then
        log_err "cuDNN installation error: CUDA directory not found, please install CUDA first"
        return 1
    fi

    # Copy files based on cuDNN version
    log_info "Starting cuDNN ${CUDNN_VERSION} installation..."
    # For cuDNN versions below 8.4, the extracted package directly contains include and lib folders.
    if [ "$CUDNN_VERSION" \< "8.4" ];then
        if [ -d "${AUTO_INSTALL}/cudnn/cuda/include" ] && [ -d "${AUTO_INSTALL}/cudnn/cuda/lib64" ]; then
            cp  ${AUTO_INSTALL}/cudnn/cuda/include/* /usr/local/cuda/include
            if [ $? -ne 0 ]; then
                log_err "cuDNN installation error: Failed to copy include files!"
                return 1
            fi
            cp -P ${AUTO_INSTALL}/cudnn/cuda/lib64/* /usr/local/cuda/lib64
            if [ $? -ne 0 ]; then
                log_err "cuDNN installation error: Failed to copy lib files!"
                return 1
            fi
        else
            log_err "cuDNN installation error: Extracted cuDNN directory structure is incorrect!"
            return 1
        fi
    else
        # Starting from cuDNN 8.4, NVIDIA changed the packaging structure. Extraction creates an archive directory with version and architecture info, need to find the directory first.
        cudnn_archive_dir=$(find ${AUTO_INSTALL}/cudnn -name "cudnn-*-archive" -type d | head -1)
        if [ -z "$cudnn_archive_dir" ]; then
            log_err "cuDNN installation error: cuDNN archive directory not found!"
            return 1
        fi
        if [ -d "$cudnn_archive_dir/include" ] && [ -d "$cudnn_archive_dir/lib" ]; then
            cp ${cudnn_archive_dir}/include/cudnn*.h /usr/local/cuda/include
            if [ $? -ne 0 ]; then
                log_err "cuDNN installation error: Failed to copy include files!"
                return 1
            fi
            cp -P ${cudnn_archive_dir}/lib/libcudnn* /usr/local/cuda/lib64
            if [ $? -ne 0 ]; then
                log_err "cuDNN installation error: Failed to copy lib files!"
                return 1
            fi
        else
            log_err "cuDNN installation error: Extracted cuDNN directory structure is incorrect!"
            return 1
        fi
    fi

    # Set file permissions
    chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*
    if [ $? -ne 0 ]; then
        log_err "cuDNN installation error: Failed to set permissions!"
        return 1
    fi

    # Verify cuDNN installation
    if [ -f /usr/local/cuda/include/cudnn_version.h ]; then
        cudnn_installed_version=$(grep "#define CUDNN_MAJOR" /usr/local/cuda/include/cudnn_version.h | awk '{print $3}').$(grep "#define CUDNN_MINOR" /usr/local/cuda/include/cudnn_version.h | awk '{print $3}').$(grep "#define CUDNN_PATCHLEVEL" /usr/local/cuda/include/cudnn_version.h | awk '{print $3}')
        log_info "cuDNN installed successfully, version: $cudnn_installed_version"
    else
        log_warn "Unable to verify cuDNN version, but installation process completed"
    fi

    log_always "cuDNN $CUDNN_VERSION installed successfully!"
    return 0
}

enable_pm() {
    if systemctl list-units --type=service | grep -q "nvidia-persistenced.service"; then
        log_info "nvidia-persistenced.service already exists, skipping..."
        return 0
    fi

    if command -v bunzip2 >/dev/null 2>&1; then
        log_debug "bunzip2 command is available"

        # Check if sample directory exists
        if [ ! -d "/usr/share/doc/NVIDIA_GLX-1.0" ]; then
            log_err "NVIDIA sample directory does not exist, cannot enable persistence mode"
            return 1
        fi

        cd /usr/share/doc/NVIDIA_GLX-1.0/sample* || {
            log_err "Unable to switch to sample directory"
            return 1
        }

        # Check if compressed file exists
        if [ ! -f "nvidia-persistenced-init.tar.bz2" ]; then
            log_err "nvidia-persistenced-init.tar.bz2 file does not exist"
            return 1
        fi

        # Extract file
        tar -xvjf nvidia-persistenced-init.tar.bz2 || {
            log_err "Failed to extract nvidia-persistenced-init.tar.bz2"
            return 1
        }

        # Install persistence service
        cd nvidia-persistenced-init && sh install.sh -u root || {
            log_err "Failed to install nvidia-persistenced service"
            return 1
        }

        log_info "GPU persistence mode configured successfully"
    else
        log_err "bunzip2 is not installed, need to manually enable persistence mode later!"
        return 1
    fi

    return 0
}

disable_GSP_module() {

    sh -c 'echo "options nvidia NVreg_EnableGpuFirmware=0" > /etc/modprobe.d/nvidia-gsp.conf'

    case "$os" in
        ubuntu)
            log_debug "Detected Ubuntu system"
            update-initramfs -u
            ;;
        centos)
            log_debug "Detected CentOS system"
            dracut -f
            ;;
        ctyunos)
            log_debug "Detected CTyunOS system"
            dracut -f
            ;;
        kylin)
            log_debug "Detected Kylin system"
            dracut -f
            ;;
        *)
            log_err "Invalid operating system! disable_GSP_module failed! os is $os"
            return 1
            ;;
    esac

    log_info "GSP module disable configured successfully, will take effect after reboot!"
}

# Validate YUM source connectivity
validate_yum_source() {
    # 1. Parameter check (required validation)
    if [ $# -ne 1 ] || [ -z "$1" ]; then
        log_err "validate_yum_source function call error: must pass 1 parameter (YUM source address)"
        return 1
    fi

    local yum_source=$1  # Local variable to avoid global pollution
    local max_retry=6   # Maximum retry count
    local retry_interval=5  # Retry interval (seconds)
    local status_code

    log_info "Starting YUM source address validation: $yum_source"

    # Try primary source with max retries
    local yum_source_ok=false
    for ((i=1; i<=max_retry; i++))
    do
        status_code=$(curl -s -o /dev/null -w "%{http_code}" "${yum_source}")
        if [ "$status_code" = "200" ]; then
            log_info "Primary YUM repository Access granted: $yum_source"
            yum_source_ok=true
            break
        else
            log_warn "Primary YUM repository Access denied, HTTP code: $status_code (attempt $i/$max_retry)"
            sleep $retry_interval
        fi
    done

    # If primary source failed, try backup source with 12 retries
    if [ "$yum_source_ok" = "false" ]; then
        log_warn "Primary YUM source failed after $max_retry attempts, trying backup source: $yum_source_back"
        
        for ((i=1; i<=max_retry; i++))
        do
            status_code=$(curl -s -o /dev/null -w "%{http_code}" "${yum_source_back}")
            if [ "$status_code" = "200" ]; then
                log_info "Backup YUM repository Access granted: $yum_source_back (attempt $i/$max_retry)"
                yum_source=$yum_source_back
                log_info "Using backup source as primary: $yum_source"
                yum_source_ok=true
                break
            else
                log_warn "Backup YUM repository Access denied, HTTP code: $status_code (attempt $i/$max_retry)"
                sleep $retry_interval
            fi
        done
        
        if [ "$yum_source_ok" = "false" ]; then
            log_err "Backup YUM repository also failed after $max_retry attempts, HTTP code: $status_code"
            log_err "Both YUM repositories are unreachable, installation cannot proceed"
            return 1
        fi
    fi
    
    return 0
}

# Get system and configuration file information
get_system_and_env_info() {
    local os=""
    local profile_file=""
    local env_file=""

    # 1. Prefer getting system identifier from /etc/os-release (standard method)
    if [ -f "/etc/os-release" ]; then
        os=$(. /etc/os-release; echo "$ID")
        os=$(echo "$os" | tr '[:upper:]' '[:lower:]')  # Convert to lowercase to avoid case issues

        # 2. Match configuration file paths based on system type
        case "$os" in
            "ubuntu")
                profile_file="/root/.profile"
                env_file="/root/.bashrc"
                ;;
            "centos"|"ctyunos"|"kylin")
                profile_file="/root/.bash_profile"
                env_file="/root/.bashrc"
                ;;
            *)
                # Unsupported system, return failure
                echo ""
                return 1
                ;;
        esac
    else
        # 3. Compatibility for older systems: identify CentOS from /etc/issue
        if grep -q "CentOS" /etc/issue 2>/dev/null; then
            os="centos"
            profile_file="/root/.bash_profile"
            env_file="/root/.bashrc"
        else
            # Unable to identify system, return failure
            echo ""
            return 1
        fi
    fi

    # 4. Return success: concatenate in os:profile:env format (key modification point)
    echo "${os}:${profile_file}:${env_file}"
    return 0
}

set_env() {
    # Check if CUDA installation directory exists
    if [ ! -d "/usr/local/cuda/bin" ]; then
        log_err "CUDA installation directory does not exist, cannot set environment variables"
        return 1
    fi

    env_path="/usr/local/cuda/bin:"
    env_library="/usr/local/cuda/lib64:"
    env1="export PATH=${env_path}\$PATH"
    env2="export LD_LIBRARY_PATH=${env_library}\$LD_LIBRARY_PATH"

    # Check if environment file exists
    if [ ! -f "${env_file}" ]; then
        log_warn "Environment file ${env_file} does not exist, will create new file"
        touch "${env_file}"
        if [ $? -ne 0 ]; then
            log_err "Failed to create environment file"
            return 1
        fi
    fi

    # Check if environment variable already exists
    if grep -q "export PATH=${env_path}" "${env_file}"; then
        log_info "PATH environment variable already exists, skipping"
    else
        echo $env1 >> ${env_file}
        if [ $? -ne 0 ]; then
            log_err "Failed to add PATH environment variable"
            return 1
        fi
    fi

    if grep -q "export LD_LIBRARY_PATH=${env_library}" "${env_file}"; then
        log_info "LD_LIBRARY_PATH environment variable already exists, skipping"
    else
        echo $env2 >> ${env_file}
        if [ $? -ne 0 ]; then
            log_err "Failed to add LD_LIBRARY_PATH environment variable"
            return 1
        fi
    fi

    # Verify file write success
    if [ ! -s "${env_file}" ]; then
        log_err "Environment file is empty, setup failed"
        return 1
    fi

    # Reload environment variables
    source ${env_file}
    if [ $? -ne 0 ]; then
        log_warn "Issue occurred while reloading environment variables, but installation may still be successful"
    fi

    log_info "Environment variables set successfully!"
    return 0
}

verify_installation() {
    log_info "Starting installation verification..."

    local all_ok=0

    # Verify NVIDIA driver
    log_info "Verifying NVIDIA driver..."
    if command -v nvidia-smi &> /dev/null; then
        log_info "NVIDIA driver installed, version: $(nvidia-smi --query-gpu=driver_version --format=csv,noheader)"
        # Check if GPU is correctly detected
        local gpu_count=$(nvidia-smi --query-gpu=count --format=csv,noheader)
        if [ -n "$gpu_count" ] && [ "$gpu_count" -gt 0 ]; then
            log_info "Successfully detected $gpu_count GPU(s)"
            # Display GPU info
            nvidia-smi | head -20 | tee -a "$LOG_FILE"
        else
            log_warn "No GPU detected, driver may not be loaded correctly"
            all_ok=1
        fi
    else
        log_err "NVIDIA driver not installed successfully"
        all_ok=1
    fi

    # Verify CUDA
    log_info "Verifying CUDA..."
    if command -v /usr/local/cuda/bin/nvcc &> /dev/null; then
        log_info "CUDA installed, version: $(/usr/local/cuda/bin/nvcc --version | grep release | awk '{print $NF}')"
        # Verify CUDA library exists
        if [ -f "/usr/local/cuda/lib64/libcudart.so" ]; then
            log_info "CUDA runtime library exists"
        else
            log_warn "CUDA runtime library does not exist"
            all_ok=1
        fi
    else
        log_err "CUDA not installed successfully"
        all_ok=1
    fi

    # Verify cuDNN
    log_info "Verifying cuDNN..."
    if [ -f "/usr/local/cuda/include/cudnn_version.h" ]; then
        local cudnn_version=$(grep "#define CUDNN_MAJOR" /usr/local/cuda/include/cudnn_version.h | awk '{print $3}').$(grep "#define CUDNN_MINOR" /usr/local/cuda/include/cudnn_version.h | awk '{print $3}').$(grep "#define CUDNN_PATCHLEVEL" /usr/local/cuda/include/cudnn_version.h | awk '{print $3}')
        log_info "cuDNN installed, version: $cudnn_version"
        # Verify cuDNN library exists
        if [ -f "/usr/local/cuda/lib64/libcudnn.so" ]; then
            log_info "cuDNN library exists"
        else
            log_warn "cuDNN library does not exist"
            all_ok=1
        fi
    else
        log_err "cuDNN not installed successfully"
        all_ok=1
    fi

    # Verify environment variables
    log_info "Verifying environment variables..."
    if [ -f "${env_file}" ] && grep -q "/usr/local/cuda/bin" "${env_file}" 2>/dev/null; then
        log_info "CUDA PATH environment variable configured in ${env_file}"
    else
        log_warn "CUDA PATH environment variable not configured in ${env_file}"
        all_ok=1
    fi

    if [ -f "${env_file}" ] && grep -q "/usr/local/cuda/lib64" "${env_file}" 2>/dev/null; then
        log_info "CUDA LD_LIBRARY_PATH configured in ${env_file}"
    else
        log_warn "CUDA LD_LIBRARY_PATH not configured in ${env_file}"
        all_ok=1
    fi

    # Summarize verification results
    if [ "$all_ok" -eq 0 ]; then
        log_always "All components verified successfully, installation complete!"
        return 0
    else
        log_warn "Some components failed verification, recommend checking logs and rebooting system for re-verification"
        return 1
    fi
}

check_installation() {
    if [ "$(id -u)" != "0" ]; then
        echo "[ERROR] This command must be run with root privileges"
        exit 1
    fi
    
    local driver_installed="NO"
    local cuda_installed="NO"
    local cudnn_installed="NO"
    local driver_version=""
    local cuda_version=""
    local cudnn_version=""
    
    # Check driver
    if command -v nvidia-smi &> /dev/null; then
        driver_installed="YES"
        driver_version=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1)
        if [ -z "$driver_version" ]; then
            driver_version=$(nvidia-smi 2>/dev/null | grep "Driver Version" | awk '{print $6}')
        fi
    fi
    
    # Check CUDA
    if command -v /usr/local/cuda/bin/nvcc &> /dev/null; then
        cuda_installed="YES"
        cuda_version=$(/usr/local/cuda/bin/nvcc --version 2>/dev/null | grep "release" | awk '{print $NF}' | sed 's/,//')
    fi
    
    # Check cuDNN
    if [ -f /usr/local/cuda/include/cudnn_version.h ]; then
        cudnn_installed="YES"
        cudnn_version=$(grep "#define CUDNN_MAJOR" /usr/local/cuda/include/cudnn_version.h 2>/dev/null | awk '{print $3}').$(grep "#define CUDNN_MINOR" /usr/local/cuda/include/cudnn_version.h 2>/dev/null | awk '{print $3}').$(grep "#define CUDNN_PATCHLEVEL" /usr/local/cuda/include/cudnn_version.h 2>/dev/null | awk '{print $3}')
    fi
    
    # Summary
    echo ""
    echo "=========================================="
    echo "           Installation Summary           "
    echo "=========================================="
    printf "%-10s %s\n" "Driver:" "${driver_version:-Not installed}"
    printf "%-10s %s\n" "CUDA:" "${cuda_version:-Not installed}"
    printf "%-10s %s\n" "cuDNN:" "${cudnn_version:-Not installed}"
    echo "=========================================="
    echo ""
    
    if [ "$driver_installed" = "YES" ] && [ "$cuda_installed" = "YES" ] && [ "$cudnn_installed" = "YES" ]; then
        echo "[OK] All components installed successfully!"
        return 0
    else
        echo "[FAIL] Installation incomplete!"
        return 1
    fi
}

# ===================== Main Execution Flow =====================
function main() {
    # Parse command line arguments first
    parse_arguments "$@"
    
    # If check mode, only check installation status and exit
    if [ "$CHECK_MODE" = true ]; then
        check_installation
        exit $?
    fi
    
    # Validate driver version before requiring root privileges
    if [ -n "$DRIVER_OVERRIDE" ]; then
        if ! validate_driver_version "$DRIVER_OVERRIDE"; then
            exit 1
        fi
    fi
    
    check_root
    init_dir
    check_and_install_commands "gcc" "make" "wget" "tar" "curl" "bzip2"
    if [ $? -ne 0 ]; then
        log_err "Install error: Check and install commands failed!"
        exit 1
    fi

    if validate_yum_source "${yum_source}"; then
        log_info "YUM source verified, continuing with installation process..."
        # Logic after successful verification (e.g., install dependencies, configure YUM source, etc.)
    else
        log_err "YUM source verification failed, script exiting"
        exit 1
    fi

    sys_info=$(get_system_and_env_info)

    # 2. Check if function executed successfully
    if [ -z "$sys_info" ]; then
        log_err "Error: Unable to identify current operating system, only ubuntu/centos/ctyunos/kylin are supported"
        exit 1
    fi

    # 3. Split result into os, profile_file, env_file (colon-separated)
    os=$(echo "$sys_info" | cut -d':' -f1)
    profile_file=$(echo "$sys_info" | cut -d':' -f2)
    env_file=$(echo "$sys_info" | cut -d':' -f3)
    log_always "os:$os"

    select_version

   # Start download
    begin_download=$(date '+%s')

    download
    if [ $? -ne 0 ]; then
        log_err "Install error: Download failed!"
        exit 1
    fi
    log_info "Download completed!"

    end_download=$(date '+%s')
    time_download=$((end_download-begin_download))
    log_info "NVIDIA download completed! Time elapsed: $time_download seconds"

    disable_nouveau

    begin=$(date '+%s')

    install_driver
    if [ $? -ne 0 ]; then
        log_err "Install error: Driver installation failed!"
        exit 1
    fi

    install_cuda
    if [ $? -ne 0 ]; then
        log_err "Install error: CUDA installation failed!"
        exit 1
    fi

    install_cudnn
    if [ $? -ne 0 ]; then
        log_err "Install error: cuDNN installation failed!"
        exit 1
    fi

    enable_pm

    set_env

    end=$(date '+%s')
    time_install=$((end-begin))

    log_info "------------------------------------------------"
    log_info "All components installed successfully"
    log_info "Total installation time: $time_install seconds"
    log_info "Installation completed."
    log_info "------------------------------------------------"

    # Run post-installation verification
    verify_installation
    if [ $? -ne 0 ]; then
        log_warn "Installation verification found issues, but system will continue to reboot to apply changes"
    else
        log_info "Installation verified successfully, system will reboot to complete configuration"
    fi

    # Update dynamic linker cache
    ldconfig
    log_info "Rebooting system..."
    sleep 30
    reboot
}

main "$@"
