#!/bin/bash
# =============================================================================
# loonggpu-ctk - LoongGPU CDI(Container Device Interface) 配置工具
#
# 为玲珑(ll-cli)容器注入 LoongGPU 设备与用户态驱动。
#
#   loonggpu-ctk generate   根据当前 GPU 设备生成/刷新 CDI 规范与玲珑配置(默认动作)
#   loonggpu-ctk clean      删除本工具生成的全部文件
#   loonggpu-ctk --help
#
# 生成:
#   /etc/cdi/<driver>.yaml               CDI 规范(设备节点 + 驱动挂载)
#   /etc/linglong/config.d/<driver>.json 玲珑使能(devices: ["<kind>=all"])
# 配置:
#   /etc/loonggpu-ctk/drivers.d/<driver>.conf
# =============================================================================
set -euo pipefail

CONFIG_DIR="/etc/loonggpu-ctk/drivers.d"
CDI_OUTPUT_DIR="/etc/cdi"
LINGLONG_OUTPUT_DIR="/etc/linglong/config.d"
CDI_VERSION="0.5.0"
MOUNT_OPTS=("ro" "nosuid" "nodev" "rbind" "rprivate")

usage() {
    cat <<EOF
用法: loonggpu-ctk <命令>

命令:
  generate    根据当前 GPU 设备生成/刷新 CDI 规范与玲珑配置(默认)
  clean       删除由本工具生成的全部文件
  --help,-h   显示帮助
EOF
}

# 按主机架构探测库目录, 导出 LIB_DIR 供驱动配置引用
detect_lib_dir() {
    case "$(uname -m)" in
        loongarch64) [ -d /usr/lib/loongarch64-linux-gnu ] && echo /usr/lib/loongarch64-linux-gnu || { [ -d /usr/lib64 ] && echo /usr/lib64 || echo /usr/lib; } ;;
        x86_64)      [ -d /usr/lib/x86_64-linux-gnu ]      && echo /usr/lib/x86_64-linux-gnu      || { [ -d /usr/lib64 ] && echo /usr/lib64 || echo /usr/lib; } ;;
        aarch64)     [ -d /usr/lib/aarch64-linux-gnu ]     && echo /usr/lib/aarch64-linux-gnu     || { [ -d /usr/lib64 ] && echo /usr/lib64 || echo /usr/lib; } ;;
        *)           echo /usr/lib ;;
    esac
}
LIB_DIR=$(detect_lib_dir)

# 原子写: 内容来自 stdin, 写入同目录临时文件后 rename, 属主 root:root 0644
write_atomic() {
    local dest=$1 tmp
    tmp=$(mktemp "${dest}.tmp.XXXXXX") || return 1
    cat > "$tmp"
    chmod 644 "$tmp"
    chown root:root "$tmp" 2>/dev/null || true
    mv -f "$tmp" "$dest"
}

# 扫描绑定指定内核驱动的 DRM 设备。输出: card名|render名(逗号分隔)|排序键
scan_gpus() {
    local driver_name=$1
    local -n gpu_arr=$2
    gpu_arr=()
    local drm_path="/sys/class/drm"
    [ ! -d "$drm_path" ] && return

    local card_sysfs
    for card_sysfs in "$drm_path"/card[0-9]*; do
        [ -e "$card_sysfs" ] || continue
        local card_name driver_link current_driver
        card_name=$(basename "$card_sysfs")
        driver_link="$card_sysfs/device/driver"
        [ ! -L "$driver_link" ] && continue
        current_driver=$(basename "$(readlink -f "$driver_link")")
        [ "$current_driver" != "$driver_name" ] && continue

        local render_names="" drm_dev_dir="$card_sysfs/device/drm" dev
        if [ -d "$drm_dev_dir" ]; then
            for dev in "$drm_dev_dir"/renderD*; do
                [ -e "$dev" ] || continue
                [ -n "$render_names" ] && render_names+=","
                render_names+=$(basename "$dev")
            done
        fi

        local sort_key="$card_sysfs" resolved
        if [ -L "$card_sysfs/device" ]; then
            resolved=$(readlink -f "$card_sysfs/device")
            [ -n "$resolved" ] && sort_key=$(basename "$resolved")
        fi
        gpu_arr+=("${card_name}|${render_names}|${sort_key}")
    done

    if [ ${#gpu_arr[@]} -gt 0 ]; then
        local sorted
        IFS=$'\n' sorted=($(sort -t'|' -k3 <<<"${gpu_arr[*]}"))
        unset IFS
        gpu_arr=("${sorted[@]}")
    fi
}

# 输出一个挂载项(仅当宿主路径存在)
emit_mount_if_exists() {
    local host_path=$1
    if [ -e "$host_path" ] || [ -L "$host_path" ]; then
        echo "        - hostPath: \"${host_path}\""
        echo "          containerPath: \"${host_path}\""
        echo "          options:"
        local opt
        for opt in "${MOUNT_OPTS[@]}"; do echo "            - \"${opt}\""; done
    else
        echo "WARN: 跳过不存在的主机路径: $host_path" >&2
    fi
}

emit_mounts() {
    local driver_name=$1 pattern file
    local patterns=(${MOUNT_FILES[@]+"${MOUNT_FILES[@]}"})
    if [ ${#patterns[@]} -eq 0 ]; then
        echo "WARN: 驱动 [$driver_name] 未定义 MOUNT_FILES, 不生成挂载" >&2
        return
    fi
    shopt -s nullglob
    for pattern in "${patterns[@]}"; do
        [ -z "$pattern" ] && continue
        for file in $pattern; do emit_mount_if_exists "$file"; done
    done
    shopt -u nullglob
}

remove_generated() { # $1=driver_name
    local d=$1 f
    for f in "${CDI_OUTPUT_DIR}/${d}.yaml" "${LINGLONG_OUTPUT_DIR}/${d}.json"; do
        if [ -e "$f" ] || [ -L "$f" ]; then echo "Removing $f"; rm -f "$f"; fi
    done
}

# 处理单份驱动配置
generate_one() {
    local conf_file=$1
    unset DRIVER_NAME CDI_KIND MOUNT_FILES
    # shellcheck disable=SC1090
    source "$conf_file"

    if [ -z "${DRIVER_NAME:-}" ] || [ -z "${CDI_KIND:-}" ]; then
        echo "ERROR: 无效配置(缺 DRIVER_NAME/CDI_KIND): $conf_file" >&2
        return 1
    fi
    if ! [[ "$DRIVER_NAME" =~ ^[A-Za-z0-9_-]+$ ]]; then
        echo "ERROR: DRIVER_NAME 含非法字符: $DRIVER_NAME" >&2; return 1
    fi
    if ! [[ "$CDI_KIND" =~ ^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$ ]]; then
        echo "ERROR: CDI_KIND 非法: $CDI_KIND" >&2; return 1
    fi
    if ! declare -p MOUNT_FILES &>/dev/null; then MOUNT_FILES=(); fi

    local cdi_file="${CDI_OUTPUT_DIR}/${DRIVER_NAME}.yaml"
    local ll_file="${LINGLONG_OUTPUT_DIR}/${DRIVER_NAME}.json"

    local gpus=()
    scan_gpus "$DRIVER_NAME" gpus
    if [ ${#gpus[@]} -eq 0 ]; then
        echo "INFO: 内核驱动 [$DRIVER_NAME] 无匹配 GPU, 移除旧配置(若有)"
        remove_generated "$DRIVER_NAME"
        return 0
    fi
    echo "INFO: 驱动 [$DRIVER_NAME] 发现 ${#gpus[@]} 个 GPU"

    local info card_name render_names
    {
        echo "cdiVersion: \"${CDI_VERSION}\""
        echo "kind: \"${CDI_KIND}\""
        echo ""
        echo "devices:"
        echo "  - name: \"all\""
        echo "    containerEdits:"
        echo "      deviceNodes:"
        for info in "${gpus[@]}"; do
            IFS='|' read -r card_name render_names _ <<< "$info"
            if [ -e "/dev/dri/${card_name}" ]; then
                echo "        - hostPath: \"/dev/dri/${card_name}\""
                echo "          path: \"/dev/dri/${card_name}\""
                echo "          permissions: \"rwm\""
            fi
            local rn
            if [ -n "$render_names" ]; then
                IFS=',' read -ra rl <<< "$render_names"
                for rn in "${rl[@]}"; do
                    if [ -e "/dev/dri/${rn}" ]; then
                        echo "        - hostPath: \"/dev/dri/${rn}\""
                        echo "          path: \"/dev/dri/${rn}\""
                        echo "          permissions: \"rwm\""
                    fi
                done
            fi
        done
        echo "      mounts:"
        emit_mounts "$DRIVER_NAME"
    } | write_atomic "$cdi_file"
    echo "INFO:   生成 $cdi_file"

    {
        echo "{"
        echo "    \"devices\": ["
        echo "        \"${CDI_KIND}=all\""
        echo "    ]"
        echo "}"
    } | write_atomic "$ll_file"
    echo "INFO:   生成 $ll_file"
}

do_generate() {
    mkdir -p "$CDI_OUTPUT_DIR" "$LINGLONG_OUTPUT_DIR"
    if [ ! -d "$CONFIG_DIR" ]; then
        echo "INFO: 配置目录 $CONFIG_DIR 不存在, 退出"
        exit 0
    fi
    local count=0 failures=0 conf
    for conf in "$CONFIG_DIR"/*.conf; do
        [ -e "$conf" ] || continue
        if ( generate_one "$conf" ); then count=$((count+1)); else
            echo "ERROR: 处理失败: $conf" >&2; failures=$((failures+1));
        fi
    done
    echo "INFO: 成功处理 $count 份驱动配置"
    [ "$failures" -gt 0 ] && { echo "ERROR: $failures 份失败" >&2; return 1; }
    return 0
}

do_clean() {
    echo "INFO: 清理本工具生成的文件..."
    local conf driver_name
    for conf in "$CONFIG_DIR"/*.conf; do
        [ -e "$conf" ] || continue
        driver_name=$(sed -n 's/^[[:space:]]*DRIVER_NAME=["'"'"']\([^"'"'"']*\)["'"'"'].*/\1/p' "$conf" | head -n1)
        [ -n "$driver_name" ] && remove_generated "$driver_name"
    done
    echo "INFO: 清理完成"
}

case "${1:-generate}" in
    generate)      do_generate ;;
    clean)         do_clean ;;
    -h|--help|help) usage ;;
    *) echo "未知命令: $1" >&2; usage >&2; exit 2 ;;
esac
