#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

install_dir=$1
script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
temp_folders=()
checked_pkgs=()
global_temp_dir=""
cached_installed_pkgs=()

# Configuration
MAX_PARALLEL_JOBS=4
BATCH_SIZE=8

echo "Current system architecture: $NSYS_ARCH" >/dev/tty
echo "Using batch processing with $MAX_PARALLEL_JOBS parallel jobs" >/dev/tty

mkdir -p "$install_dir"

trap cleanup_temp_files INT

function cleanup_temp_files() {
    for temp_folder in "${temp_folders[@]}"; do
        rm -rf "$temp_folder"
    done
    temp_folders=()
    if [ -n "$global_temp_dir" ]; then
        rm -rf "$global_temp_dir"
    fi
}

function contains_element() {
    local e match=$1
    shift
    for e; do [[ "$e" == "$match" ]] && return 0; done
    return 1
}

function replace_in_file() {
    local old=$1
    local new=$2
    local file_to_replace=$3
    sed -i -e "s@$old@$new@g" "$file_to_replace"
}

function replace_abs_path_in_file() {
    local path=$1
    local file_to_replace=$2
    replace_in_file "$path" "$install_dir$path" "$file_to_replace"
}

# force cache update with our installed dir
function update_fonts_cache() {
    if [ -e "$install_dir/etc/fonts/fonts.conf" ]; then
        if [ -e "$install_dir/usr/bin/fc-cache" ]; then
            local conf_path
            conf_path=$install_dir/etc/fonts/fonts.conf
            replace_in_file "conf.d" "$install_dir/etc/fonts/conf.d" "$conf_path"
            replace_abs_path_in_file "/usr/share/fonts" "$conf_path"
            "$install_dir/usr/bin/fc-cache" -sfv
        fi
    fi
}

function get_package_manager() {
    local package_manager
    package_manager=$(which dnf 2>/dev/null)
    if [ -z "$package_manager" ]; then
        package_manager=$(which yum)
        if [ -z "$package_manager" ]; then
            echo "yum and dnf package managers are no exist. To continue, you need to install one of them." >/dev/tty
        fi
    fi
    echo "$package_manager"
}

function is_package_installed_global() {
    local pkg_name=$1
    # Use cached results to avoid repeated system queries
    for full_pkg_name in "${cached_installed_pkgs[@]}"; do
        if [[ "$full_pkg_name" == "${pkg_name}"* ]]; then
            echo "Package is installed (cached): $pkg_name" >/dev/tty
            return 0
        fi
    done
    return 1
}

function cache_installed_packages() {
    echo "Caching installed packages for faster lookups..." >/dev/tty
    local package_manager
    package_manager=$(get_package_manager)
    if [ -n "$package_manager" ]; then
        if [[ "$package_manager" == *dnf ]]; then
            mapfile -t cached_installed_pkgs < <($package_manager list installed 2>/dev/null | awk 'NR>1 {print $1}' | cut -d'.' -f1)
        else
            mapfile -t cached_installed_pkgs < <($package_manager list installed 2>/dev/null | awk 'NR>1 {print $1}' | cut -d'.' -f1)
        fi
    fi
    echo "Cached ${#cached_installed_pkgs[@]} installed packages" >/dev/tty
}

function check_install_yumutils() {
    local repoquery
    repoquery=$(which repoquery 2>/dev/null)
    if [ -z "$repoquery" ]; then
        echo "Installing yum-utils and which..." >/dev/tty
        if [ -z "$global_temp_dir" ]; then
            global_temp_dir="$install_dir/$(tr -dc 'a-zA-Z0-9' <'/dev/urandom' | fold -w 16 | head -n 1)"
            mkdir -p "$global_temp_dir"
            temp_folders+=("$global_temp_dir")
        fi

        mapfile -t utility_packages < <("$script_dir/dependency_resolver" "$global_temp_dir" "which,yum-utils" false "${checked_pkgs[@]}")
        if ((${#utility_packages[@]})); then
            readarray -t utility_installed < <("$script_dir"/rpm_extract "$install_dir" "${utility_packages[@]}")
        fi

        repoquery=$(which repoquery 2>/dev/null)
    fi
    echo "$repoquery"
}

function batch_download_packages() {
    local packages=("$@")

    echo "Batch downloading ${#packages[@]} packages..." >/dev/tty

    if [ -z "$global_temp_dir" ]; then
        global_temp_dir="$install_dir/$(tr -dc 'a-zA-Z0-9' <'/dev/urandom' | fold -w 16 | head -n 1)"
        mkdir -p "$global_temp_dir"
        temp_folders+=("$global_temp_dir")
    fi

    # Check if packages are library file names (contain .so) - if so, process individually
    local has_libraries=false
    for pkg in "${packages[@]}"; do
        if [[ "$pkg" == *".so"* ]]; then
            has_libraries=true
            break
        fi
    done

    if [ "$has_libraries" = true ]; then
        echo "Detected library file names, processing with parallel provider resolution..." >/dev/tty
        local all_rpm_files=()
        local job_count=0
        local pids=()

        echo "Processing ${#packages[@]} individual libraries..." >/dev/tty
        for pkg_name in "${packages[@]}"; do
            echo "DEBUG: Processing individual library: '$pkg_name'" >/dev/tty
            if [ $job_count -ge $MAX_PARALLEL_JOBS ]; then
                # Wait for one job to complete
                wait "${pids[0]}"
                pids=("${pids[@]:1}")  # Remove first element
                ((job_count--))
            fi

            # Start background job for library resolution
            {
                echo "Resolving provider for individual library: $pkg_name" >/dev/tty
                mapfile -t pkg_rpms < <("$script_dir/dependency_resolver" "$global_temp_dir" "$pkg_name" false "${checked_pkgs[@]}")

                if ((${#pkg_rpms[@]})); then
                    # Write RPMs to a temp file named after the process
                    local temp_file="$global_temp_dir/rpms_$$_${BASHPID}.list"
                    printf '%s\n' "${pkg_rpms[@]}" > "$temp_file"
                    echo "Found ${#pkg_rpms[@]} RPMs for $pkg_name" >/dev/tty
                else
                    echo "No RPMs found for $pkg_name" >/dev/tty
                fi
            } &

            pids+=($!)
            ((job_count++))
        done

        for pid in "${pids[@]}"; do
            wait "$pid"
        done

        # Collect all RPM files from temp files
        shopt -s nullglob
        local temp_rpm_files=("$global_temp_dir"/rpms_*.list)
        shopt -u nullglob

        if ((${#temp_rpm_files[@]})); then
            # Combine all RPM files and remove duplicates
            cat "${temp_rpm_files[@]}" | sort -u > "$global_temp_dir/batch_rpms.list"
            mapfile -t unique_rpms < "$global_temp_dir/batch_rpms.list"

            if ((${#unique_rpms[@]})); then
                echo "batch:${unique_rpms[*]}" > "$global_temp_dir/batch.files"
                echo "Downloaded ${#unique_rpms[@]} unique RPM files from ${#packages[@]} libraries" >/dev/tty
            fi

            # Cleanup temp files
            rm -f "${temp_rpm_files[@]}"
        else
            echo "No RPMs found for any libraries" >/dev/tty
        fi
    else
        local package_string
        package_string=$(IFS=','; echo "${packages[*]}")

        echo "Downloading regular packages in single batch: $package_string" >/dev/tty
        mapfile -t dep_pkgs < <("$script_dir/dependency_resolver" "$global_temp_dir" "$package_string" false "${checked_pkgs[@]}")

        if ((${#dep_pkgs[@]})); then
            shopt -s nullglob
            local rpm_files
            rpm_files=("$global_temp_dir"/*.rpm)
            shopt -u nullglob
            if ((${#rpm_files[@]})); then
                echo "batch:${rpm_files[*]}" > "$global_temp_dir/batch.files"
                echo "Downloaded ${#rpm_files[@]} RPM files in batch" >/dev/tty
            fi
        fi
    fi

    echo "Batch download completed" >/dev/tty
}

function batch_extract_packages() {
    local packages=("$@")

    echo "Batch extracting packages..." >/dev/tty

    # Check for batch files
    local all_rpm_files=()
    if [ -f "$global_temp_dir/batch.files" ]; then
        local batch_info
        batch_info=$(cat "$global_temp_dir/batch.files")
        local rpm_files
        rpm_files=${batch_info#*:}
        all_rpm_files+=($rpm_files)
    else
        for pkg_name in "${packages[@]}"; do
            if [ -f "$global_temp_dir/$pkg_name.files" ]; then
                local pkg_info
                pkg_info=$(cat "$global_temp_dir/$pkg_name.files")
                local rpm_files
                rpm_files=${pkg_info#*:}
                all_rpm_files+=($rpm_files)
            fi
        done
    fi

    if ((${#all_rpm_files[@]})); then
        echo "Extracting ${#all_rpm_files[@]} RPM files..." >/dev/tty
        readarray -t installed < <("$script_dir"/rpm_extract "$install_dir" "${all_rpm_files[@]}")
        echo "Extracted ${#installed[@]} packages" >/dev/tty

        for pkg_name in "${packages[@]}"; do
            checked_pkgs+=("$pkg_name")
        done
    fi
}

function get_packages_to_install() {
    local all_packages=("$@")
    local packages_to_install=()

    echo "Checking which packages need installation..." >/dev/tty

    for pkg_name in "${all_packages[@]}"; do
        if is_package_installed_global "$pkg_name"; then
            echo "Package already installed: $pkg_name" >/dev/tty
        elif contains_element "$pkg_name" "${checked_pkgs[@]}"; then
            echo "Package already processed: $pkg_name" >/dev/tty
        else
            packages_to_install+=("$pkg_name")
        fi
    done

    # Output each package on a separate line to preserve array structure
    printf '%s\n' "${packages_to_install[@]}"
}

function process_package_batch() {
    local batch_packages=("$@")

    if [ ${#batch_packages[@]} -eq 0 ]; then
        return
    fi

    echo "Processing batch of ${#batch_packages[@]} packages: ${batch_packages[*]}" >/dev/tty

    batch_download_packages "${batch_packages[@]}"
    batch_extract_packages "${batch_packages[@]}"
}

echo "Starting RHEL dependencies installer..." >/dev/tty

# Cache installed packages once at the beginning
cache_installed_packages

# Install yum-utils if needed
repoquery=$(check_install_yumutils)
if [ -n "$repoquery" ]; then
    # Cache system packages for dependency resolver
    if [[ "$(get_package_manager)" == *dnf ]]; then
        repoquery_cmd="$repoquery --arch=$NSYS_ARCH,noarch"
    else
        repoquery_cmd="$repoquery --archlist=$NSYS_ARCH,noarch"
    fi

    readarray -t installed_on_system < <($repoquery_cmd -a --installed --qf="%{name}-%{version}-%{release}.%{arch}" 2>/dev/null)
    checked_pkgs+=("${installed_on_system[@]}")

    readarray -t deps_list <"$script_dir/packages.list"
    mapfile -t packages_to_install < <(get_packages_to_install "${deps_list[@]}")

    if [ ${#packages_to_install[@]} -eq 0 ]; then
        echo "All packages are already installed!" >/dev/tty
    else
        echo "Need to install ${#packages_to_install[@]} packages" >/dev/tty

        batch=()
        batch_count=0

        for pkg_name in "${packages_to_install[@]}"; do
            batch+=("$pkg_name")
            ((batch_count++))

            if [ $batch_count -ge $BATCH_SIZE ]; then
                process_package_batch "${batch[@]}"
                batch=()
                batch_count=0
            fi
        done

        # Process remaining packages in the last batch
        if [ ${#batch[@]} -gt 0 ]; then
            process_package_batch "${batch[@]}"
        fi
    fi
fi

update_fonts_cache
cleanup_temp_files

echo "Installation completed successfully!" >/dev/tty 