#!/command/with-contenv bash
set -euo pipefail

# Common helpers as requirement
if [[ -f /etc/nestri/common.sh ]]; then
    source /etc/nestri/common.sh
else
    echo "Error: Common script not found at /etc/nestri/common.sh" >&2
    exit 1
fi

# Configuration
CACHE_DIR="${NESTRI_HOME}/.cache/nestri"
NVIDIA_INSTALLER_DIR="/tmp"


# Methods #

# Ensures user ownership across directories
handle_user_permissions() {
    if ! chown "${NESTRI_USER}:${NESTRI_USER}" "${NESTRI_HOME}" 2>/dev/null; then
        echo "Error: Failed to change ownership of ${NESTRI_HOME} to ${NESTRI_USER}:${NESTRI_USER}" >&2
        return 1
    fi
    # Also apply to .cache
    if [[ -d "${NESTRI_HOME}/.cache" ]]; then
        if ! chown "${NESTRI_USER}:${NESTRI_USER}" "${NESTRI_HOME}/.cache" 2>/dev/null; then
            echo "Error: Failed to change ownership of ${NESTRI_HOME}/.cache to ${NESTRI_USER}:${NESTRI_USER}" >&2
            return 1
        fi
    fi
    return 0
}

# Prepares environment for namespace-less applications (like Steam)
setup_namespaceless() {
    rm -f /run/systemd/container || true
    mkdir -p /run/pressure-vessel || true
}

# Ensures cache directory exists
setup_cache() {
    log "Setting up cache directory at $CACHE_DIR..."
    mkdir -p "$CACHE_DIR" || {
        log "Warning: Failed to create cache directory, continuing.."
        return 1
    }
    chown "${NESTRI_USER}:${NESTRI_USER}" "$CACHE_DIR" 2>/dev/null || {
        log "Warning: Failed to set cache directory ownership, continuing.."
    }
}

# Grabs NVIDIA driver installer
get_nvidia_installer() {
    local driver_version="$1"
    local arch="$2"
    local filename="NVIDIA-Linux-${arch}-${driver_version}.run"
    local cached_file="${CACHE_DIR}/${filename}"
    local tmp_file="${NVIDIA_INSTALLER_DIR}/${filename}"

    # Check cache
    if [[ -f "$cached_file" ]]; then
        log "Found cached NVIDIA installer at $cached_file."
        cp "$cached_file" "$tmp_file" || {
            log "Warning: Failed to copy cached installer, proceeding with download."
            rm -f "$cached_file" 2>/dev/null
        }
    fi

    # Download if not in tmp
    if [[ ! -f "$tmp_file" ]]; then
        log "Downloading NVIDIA driver installer ($filename)..."
        local urls=(
            "https://international.download.nvidia.com/XFree86/Linux-${arch}/${driver_version}/${filename}"
            "https://international.download.nvidia.com/tesla/${driver_version}/${filename}"
        )
        local success=0
        for url in "${urls[@]}"; do
            if wget -q --show-progress "$url" -O "$tmp_file"; then
                success=1
                break
            fi
            log "Failed to download from $url, trying next source..."
        done

        if [[ "$success" -eq 0 ]]; then
            log "Error: Failed to download NVIDIA driver from all sources."
            return 1
        fi

        # Cache the downloaded file
        cp "$tmp_file" "$cached_file" 2>/dev/null && \
            chown "${NESTRI_USER}:${NESTRI_USER}" "$cached_file" 2>/dev/null || \
            log "Warning: Failed to cache NVIDIA driver, continuing..."
    fi

    chmod +x "$tmp_file" || {
        log "Error: Failed to make NVIDIA installer executable."
        return 1
    }
    return 0
}

# Installs the NVIDIA driver
install_nvidia_driver() {
    local filename="$1"
    log "Installing NVIDIA driver components from $filename..."
    bash ./"$filename" \
        --silent \
        --skip-depmod \
        --skip-module-unload \
        --no-kernel-module \
        --install-compat32-libs \
        --no-nouveau-check \
        --no-nvidia-modprobe \
        --no-systemd \
        --no-rpms \
        --no-backup \
        --no-distro-scripts \
        --no-libglx-indirect \
        --no-install-libglvnd \
        --no-check-for-alternate-installs || {
        log "Error: NVIDIA driver installation failed."
        return 1
    }

    log "NVIDIA driver installation completed."
    return 0
}

log_container_info() {
    if ! declare -p container_runtime &>/dev/null; then
        log "Warning: container_runtime is not defined"
        return
    fi

    if [[ "${container_runtime:-none}" != "none" ]]; then
        log "Detected container:"
        log "> ${container_runtime}"
    else
        log "No container runtime detected"
    fi
}

log_gpu_info() {
    if ! declare -p vendor_devices &>/dev/null; then
        log "Warning: vendor_devices array is not defined"
        return
    fi

    log "Detected GPUs:"
    for vendor in "${!vendor_devices[@]}"; do
        log "> $vendor: ${vendor_devices[$vendor]}"
    done
}


# Main #

# Start by getting the container we are running under
get_container_info || {
    log "Warning: Failed to detect container information"
}
log_container_info

# Setup cache now
setup_cache

# Get and detect GPU(s)
get_gpu_info || {
    log "Error: Failed to detect GPU information"
    exit 1
}
log_gpu_info

# Handle NVIDIA GPU
if [[ -n "${vendor_devices[nvidia]:-}" ]]; then
    log "NVIDIA GPU(s) detected, applying driver fix.."

    # Determine NVIDIA driver version
    nvidia_driver_version=""
    if [[ -f "/proc/driver/nvidia/version" ]]; then
        nvidia_driver_version=$(awk '/NVIDIA/ {for(i=1;i<=NF;i++) if ($i ~ /^[0-9]+\.[0-9\.]+/) {print $i; exit}}' /proc/driver/nvidia/version | head -n1)
    elif command -v nvidia-smi >/dev/null 2>&1; then
        nvidia_driver_version=$(nvidia-smi --version | grep -i 'driver version' | cut -d: -f2 | tr -d ' ')
    fi

    if [[ -z "$nvidia_driver_version" ]]; then
        log "Error: Failed to determine NVIDIA driver version."
        # Check for other GPU vendors before exiting
        if [[ -n "${vendor_devices[amd]:-}" || -n "${vendor_devices[intel]:-}" ]]; then
            log "Other GPUs (AMD or Intel) detected, continuing without NVIDIA driver"
        else
            log "No other GPUs detected, exiting due to NVIDIA driver version failure"
            exit 1
        fi
    else
        log "Detected NVIDIA driver version: $nvidia_driver_version"

        # Get installer
        arch=$(uname -m)
        filename="NVIDIA-Linux-${arch}-${nvidia_driver_version}.run"
        cd "$NVIDIA_INSTALLER_DIR" || {
            log "Error: Failed to change to $NVIDIA_INSTALLER_DIR."
            exit 1
        }
        get_nvidia_installer "$nvidia_driver_version" "$arch" || {
            # Check for other GPU vendors before exiting
            if [[ -n "${vendor_devices[amd]:-}" || -n "${vendor_devices[intel]:-}" ]]; then
                log "Other GPUs (AMD or Intel) detected, continuing without NVIDIA driver"
            else
                log "No other GPUs detected, exiting due to NVIDIA installer failure"
                exit 1
            fi
        }

        # Install driver
        install_nvidia_driver "$filename" || {
            # Check for other GPU vendors before exiting
            if [[ -n "${vendor_devices[amd]:-}" || -n "${vendor_devices[intel]:-}" ]]; then
                log "Other GPUs (AMD or Intel) detected, continuing without NVIDIA driver"
            else
                log "No other GPUs detected, exiting due to NVIDIA driver installation failure"
                exit 1
            fi
        }
    fi
fi

# Handle user permissions
log "Ensuring user permissions..."
handle_user_permissions || exit 1

# Setup namespaceless env if needed for container runtime
if [[ "$container_runtime" != "podman" ]]; then
    log "Applying namespace-less configuration"
    setup_namespaceless
fi