#!/bin/bash
log.warn() { echo -e "[\e[33mWARN\e[0m]:  \e[1m$*\e[0m"; }
log.error()  { echo -e "[\e[31mERROR\e[0m]: \e[1m$*\e[0m"; }
log.info() { echo -e "[\e[96mINFO\e[0m]:  \e[1m$*\e[0m"; }
log.debug()  { echo -e "[\e[32mDEBUG\e[0m]: \e[1m$*\e[0m"; }

# 1\. 获取宿主机 NVIDIA 驱动版本
nvidia_version=$(cat /sys/module/nvidia/version 2>/dev/null)
if [ -z "$nvidia_version" ]; then
    log.error "无法获取 NVIDIA 驱动版本"
    exit 1
fi

# 2\. 目标目录准备
ACE_DIR="/opt/apps/amber-ce-bookworm/files/ace-env"
if [[ ! -e "${ACE_DIR}" ]];then
log.error "未检测到 ACE Bookworm 安装，请安装后再试"
spark-store spk://store/tools/cn.flamescion.bookworm-compatibility-mode &
log.info "请按回车关闭..."
read
exit 1
fi
mkdir -p "$ACE_DIR/usr/lib" "$ACE_DIR/usr/lib32"

# 3\. 收集库文件路径
lib_list=$(ldconfig -p | grep -Ei "nvidia|libcuda" | cut -d'>' -f2)

# 4\. 复制库文件
copied=0
for lib in $lib_list; do
    resolved=$(readlink -f "$lib") # 解析符号链接
    if file "$resolved" | grep -q "32-bit"; then
        cp -L "$resolved" "$ACE_DIR/usr/lib32/$(basename $lib)"
    else
        cp -L "$resolved" "$ACE_DIR/usr/lib/$(basename $lib)"
        copied=1
    fi
done

# 5\. 复制辅助文件
additional_files=(
    /usr/share/vulkan/icd.d/nvidia_icd.json
    /usr/share/egl/egl_external_platform.d/20_nvidia_xcb.json
)
for file in "${additional_files[@]}"; do
    if [ -f "$file" ]; then
        dir=$(dirname "$file")
        mkdir -p "$ACE_DIR/$dir"
        cp "$file" "$ACE_DIR/$dir"
    fi
done

# 6\. 标记版本
if [ $copied -eq 1 ]; then
    echo "$nvidia_version" > "$ACE_DIR/current_version"
    log.info "NVIDIA 驱动库已成功复制"
else
    log.error "未找到有效 NVIDIA 库文件"
fi