#!/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"; }

if [ "$UID" != "0" ];then
    log.error "需要以root权限运行 Need to be run as root."
    exit 1
fi

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

# 2. 目标目录准备
ACE_DIR="$1"
if [[ ! -e "${ACE_DIR}" ]];then
    log.error "未检测到 ace-env 安装，请安装后再试 ace-env is not detected. Please try again after installation"
    exit 1
fi

# 3. 检查容器内驱动版本是否与宿主机一致
version_file="$ACE_DIR/amber-ce-tools/nvidia_current_version"
if [ -f "$version_file" ] && [ "$(cat "$version_file")" = "$nvidia_version" ]; then
#    log.info "容器内 NVIDIA 驱动版本 ($nvidia_version) 与宿主机一致，无需重新链接"
#    log.info "Driver version in container matches host version, no need to relink"
    exit 0
fi

mkdir -p "$ACE_DIR/usr/lib" "$ACE_DIR/usr/lib32"

log.info "正在链接 NVIDIA 驱动库 Linking NVIDIA Driver Libs"

# 4. 收集库文件路径
lib_list=$(ldconfig -p | grep -Ei "nvidia|libcuda|libvdpau_nvidia|libnvcuvid|libnvencode|libnvidia-encode" | cut -d'>' -f2)

# 5. 复制库文件
copied=0
has_32bit=false
has_64bit=false
has_glx=false

for lib in $lib_list; do
    resolved=$(readlink -f "$lib") # 解析符号链接
    filename=$(basename "$lib")
    if file "$resolved" | grep -q "32-bit"; then
        ln -sf "/host/$resolved" "$ACE_DIR/usr/lib32/$filename"
        has_32bit=true
        log.debug "创建32位链接: $ACE_DIR/usr/lib32/$filename -> /host/$resolved"
    else
        ln -sf "/host/$resolved" "$ACE_DIR/usr/lib/$filename"
        has_64bit=true
        copied=1
        log.debug "创建64位链接: $ACE_DIR/usr/lib/$filename -> /host/$resolved"
    fi
    
    # 检查是否是GLX库
    if [[ "$filename" == libGLX_nvidia.so.* ]]; then
        has_glx=true
    fi
done

# 6. 复制辅助文件
additional_files=(
    /usr/share/vulkan/icd.d/nvidia_icd.json
    /usr/share/vulkan/icd.d/nvidia_icd.x86_64.json
    /usr/share/vulkan/icd.d/nvidia_icd.aarch64.json
    /usr/share/vulkan/implicit_layer.d/nvidia_layers.json
    /usr/share/egl/egl_external_platform.d/10_nvidia_wayland.json
    /usr/share/egl/egl_external_platform.d/15_nvidia_gbm.json
    /usr/share/egl/egl_external_platform.d/20_nvidia_xcb.json
    /usr/share/glvnd/egl_vendor.d/10_nvidia.json
)
for file in "${additional_files[@]}"; do
    if [ -f "$file" ]; then
        resolved=$(readlink -f "$file")
        dir=$(dirname "$file")
        mkdir -p "$ACE_DIR/$dir"
        ln -sf "/host/$resolved" "$ACE_DIR/$dir/$(basename "$file")"
        log.debug "复制配置文件: $ACE_DIR/$dir/$(basename "$file") -> /host/$resolved"
    fi
done

# 7. 生成 ld.so.conf
echo "/usr/lib" > "$ACE_DIR/etc/ld.so.conf"
if [ "$has_32bit" = true ]; then
    echo "/usr/lib32" >> "$ACE_DIR/etc/ld.so.conf"
fi

# 8. 标记版本
if [ $copied -eq 1 ]; then
    echo "$nvidia_version" > "$ACE_DIR/amber-ce-tools/nvidia_current_version"
    log.info "NVIDIA 驱动库已成功链接 Nvidia Driver Libs are successfully linked."
    log.info "驱动版本: $nvidia_version"
    log.info "64位库: $has_64bit, 32位库: $has_32bit, GLX支持: $has_glx"
else
    log.error "未找到有效 NVIDIA 库文件 No valid NVIDIA Driver Libs found."
fi