#!/bin/bash
set -euo pipefail

# === CONFIGURATION SECTION ==============================
CONFIG_DIR="/etc/gpu-cdi/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")
# =========================================================

detect_lib_dir() {
    local arch=$(uname -m)
    case "$arch" in
        x86_64)
            if [ -d "/usr/lib/x86_64-linux-gnu" ]; then
                echo "/usr/lib/x86_64-linux-gnu"
            elif [ -d "/usr/lib64" ]; then
                echo "/usr/lib64"
            else
                echo "/usr/lib"
            fi
            ;;
        aarch64)
            if [ -d "/usr/lib/aarch64-linux-gnu" ]; then
                echo "/usr/lib/aarch64-linux-gnu"
            elif [ -d "/usr/lib64" ]; then
                echo "/usr/lib64"
            else
                echo "/usr/lib"
            fi
            ;;
        loongarch64)
            if [ -d "/usr/lib/loongarch64-linux-gnu" ]; then
                echo "/usr/lib/loongarch64-linux-gnu"
            elif [ -d "/usr/lib64" ]; then
                echo "/usr/lib64"
            else
                echo "/usr/lib"
            fi
            ;;
        *)
            echo "/usr/lib"
            ;;
    esac
}
LIB_DIR=$(detect_lib_dir)

mkdir -p "$CDI_OUTPUT_DIR"
mkdir -p "$LINGLONG_OUTPUT_DIR"

# Remove generated files based on the DRIVER_NAME defined in each config
if [ "${1:-}" = "--clean" ]; then
    echo "INFO: Cleaning all generated CDI configuration files..."
    for conf in "$CONFIG_DIR"/*.conf; do
        [ -e "$conf" ] || continue
        driver_name=$(sed -n 's/^[[:space:]]*DRIVER_NAME=["'"'"']\([^"'"'"']*\)["'"'"'].*/\1/p' "$conf" | head -n1)
        if [ -n "$driver_name" ]; then
            cdi_file="${CDI_OUTPUT_DIR}/${driver_name}.yaml"
            linglong_file="${LINGLONG_OUTPUT_DIR}/${driver_name}.json"
            if [ -e "$cdi_file" ] || [ -L "$cdi_file" ]; then
                echo "Removing $cdi_file"
                rm -f "$cdi_file"
            fi
            if [ -e "$linglong_file" ] || [ -L "$linglong_file" ]; then
                echo "Removing $linglong_file"
                rm -f "$linglong_file"
            fi
        fi
    done
    exit 0
fi

# Scan DRM devices bound to a kernel driver.
# Output format: card_name|render_names|sort_key
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=$(basename "$card_sysfs")

        local driver_link="$card_sysfs/device/driver"
        [ ! -L "$driver_link" ] && continue
        local current_driver=$(basename "$(readlink -f "$driver_link")")
        [ "$current_driver" != "$driver_name" ] && continue

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

        local sort_key="$card_sysfs"
        if [ -L "$card_sysfs/device" ]; then
            local resolved
            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
        IFS=$'\n' sorted=($(sort -t'|' -k3 <<<"${gpu_arr[*]}"))
        unset IFS
        gpu_arr=("${sorted[@]}")
    fi
}

# Emit a mount entry only if the host path exists
emit_mount_if_exists() {
    local host_path="$1"
    local container_path="$2"
    if [ -e "$host_path" ] || [ -L "$host_path" ]; then
        echo "        - hostPath: \"${host_path}\""
        echo "          containerPath: \"${container_path}\""
        echo "          options:"
        for opt in "${MOUNT_OPTS[@]}"; do
            echo "            - ${opt}"
        done
    else
        echo "WARN: Skipping mount for non-existent path: $host_path" >&2
    fi
}

# Emit mounts for a driver using MOUNT_FILES from the config file
emit_mounts_for_driver() {
    local driver_name="$1"
    local patterns=(${MOUNT_FILES[@]+"${MOUNT_FILES[@]}"})

    if [ ${#patterns[@]} -eq 0 ]; then
        echo "WARN: MOUNT_FILES is not defined or empty for driver [$driver_name], no mounts will be generated" >&2
        return
    fi

    shopt -s nullglob
    local pattern
    local file
    for pattern in "${patterns[@]}"; do
        [ -z "$pattern" ] && continue
        for file in $pattern; do
            emit_mount_if_exists "$file" "$file"
        done
    done
    shopt -u nullglob
}

# Generate CDI YAML and Linglong JSON for a single driver
generate_single_config() {
    local conf_file=$1

    unset DRIVER_NAME CDI_KIND MOUNT_FILES
    source "$conf_file"

    if [ -z "${DRIVER_NAME:-}" ] || [ -z "${CDI_KIND:-}" ]; then
        echo "ERROR: Invalid config file: $conf_file (missing DRIVER_NAME or 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 linglong_file="${LINGLONG_OUTPUT_DIR}/${DRIVER_NAME}.json"

    local gpus=()
    scan_gpus "$DRIVER_NAME" gpus
    local gpu_count=${#gpus[@]}

    if [ "$gpu_count" -eq 0 ]; then
        echo "INFO: No GPU device bound to driver [$DRIVER_NAME], removing old config if any"
        if [ -e "$cdi_file" ] || [ -L "$cdi_file" ]; then
            echo "Removing $cdi_file"
            rm -f "$cdi_file"
        fi
        if [ -e "$linglong_file" ] || [ -L "$linglong_file" ]; then
            echo "Removing $linglong_file"
            rm -f "$linglong_file"
        fi
        return 0
    fi

    echo "INFO: Driver [$DRIVER_NAME] - found $gpu_count GPU(s)"
    echo "INFO:   Generating CDI config: $cdi_file"
    echo "INFO:   Generating Linglong config: $linglong_file"

    {
        echo "cdiVersion: \"${CDI_VERSION}\""
        echo "kind: \"${CDI_KIND}\""
        echo ""
        echo "devices:"
        echo "  - name: \"all\""
        echo "    containerEdits:"
        echo "      deviceNodes:"

        local info
        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

            if [ -n "$render_names" ]; then
                local render_name
                IFS=',' read -ra render_list <<< "$render_names"
                for render_name in "${render_list[@]}"; do
                    if [ -e "/dev/dri/${render_name}" ]; then
                        echo "        - hostPath: \"/dev/dri/${render_name}\""
                        echo "          path: \"/dev/dri/${render_name}\""
                        echo "          permissions: \"rwm\""
                    fi
                done
            fi
        done

        echo "      mounts:"
        emit_mounts_for_driver "$DRIVER_NAME"
    } > "$cdi_file"
    chmod 644 "$cdi_file"

    {
        echo "{"
        echo "    \"devices\": ["
        echo "        \"${CDI_KIND}=all\""
        echo "    ]"
        echo "}"
    } > "$linglong_file"
    chmod 644 "$linglong_file"
}

main() {
    if [ ! -d "$CONFIG_DIR" ]; then
        echo "INFO: Driver config directory $CONFIG_DIR not found, exiting"
        exit 0
    fi

    local count=0
    local failures=0
    local conf
    for conf in "$CONFIG_DIR"/*.conf; do
        [ -e "$conf" ] || continue
        if ( generate_single_config "$conf" ); then
            count=$((count + 1))
        else
            echo "ERROR: Failed to process config: $conf" >&2
            failures=$((failures + 1))
        fi
    done

    echo "INFO: Processed $count driver config(s) successfully"
    if [ "$failures" -gt 0 ]; then
        echo "ERROR: $failures driver config(s) failed" >&2
        return 1
    fi
}

main "$@"
