#!/bin/bash
set -eo pipefail

# === CONFIGURATION SECTION - CUSTOMIZE AS NEEDED ===========
# Directory containing per-driver configuration files
CONFIG_DIR="/etc/gf-gpu-cdi/drivers.d"
# Output directory for CDI YAML configs
CDI_OUTPUT_DIR="/etc/cdi"
# Output directory for Linglong JSON configs
LINGLONG_OUTPUT_DIR="/etc/linglong/config.d"
# CDI spec version in generate yaml file
CDI_VERSION="0.5.0"
# Mount options applied to all mount entries in generate yaml file
MOUNT_OPTS=("ro" "nosuid" "nodev" "rbind" "rprivate")
# ============================================================

# found system library dir according to arch
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)

get_group_gid() {
    getent group "$1" 2>/dev/null | cut -d: -f3 || echo "0"
}
VIDEO_GID=$(get_group_gid "video")
RENDER_GID=$(get_group_gid "render")

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

# ============ delete generated config files according to generate config files driverName ============
if [ "$1" = "--clean" ]; then
    echo "INFO: Cleaning all generated CDI configuration files..."
    for conf in "$CONFIG_DIR"/*.conf; do
        [ -e "$conf" ] || continue
        # Extract DRIVER_NAME (support both single and double quotes)
        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 all GPU devices bound to the specified driver
# Output array format: card_name:render_name:pci_addr
scan_gpus() {
    local driver_name=$1
    local -n gpu_arr=$2
    gpu_arr=()

    local drm_path="/sys/class/drm"
    [ ! -d "$drm_path" ] && return

    for card_sysfs in "$drm_path"/card[0-9]*; do
        [ -e "$card_sysfs" ] || continue
        local card_name=$(basename "$card_sysfs")

        # Filter by bound driver name (skips simpledrm, virtio, other vendors)
        local driver_link="$card_sysfs/device/driver"
        [ ! -L "$driver_link" ] && continue
        local current_driver=$(basename "$(readlink -f "$driver_link")")
        [ "$current_driver" != "$driver_name" ] && continue

        # Find matching render node
        local render_name=""
        local drm_dev_dir="$card_sysfs/device/drm"
        if [ -d "$drm_dev_dir" ]; then
            for dev in "$drm_dev_dir"/renderD*; do
                if [ -e "$dev" ]; then
                    render_name=$(basename "$dev")
                    break
                fi
            done
        fi

        # Get PCI bus address (for stable sorting)
        local pci_addr=""
        if [ -L "$card_sysfs/device" ]; then
            pci_addr=$(basename "$(readlink -f "$card_sysfs/device")")
        fi

        gpu_arr+=("${card_name}:${render_name}:${pci_addr}")
    done

    # Sort by PCI address to ensure stable logical numbering (decoupled from kernel cardX order)
    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 conf file
emit_mounts_for_driver() {
    local driver_name="$1"

    # MOUNT_FILES must be defined and non-empty
    if [ ${#MOUNT_FILES[@]} -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
    for pattern in "${MOUNT_FILES[@]}"; do
        # Skip empty patterns
        [ -z "$pattern" ] && continue
        # Expand glob if any; if no match, loop is skipped
        for file in $pattern; do
            emit_mount_if_exists "$file" "$file"
        done
    done
    shopt -u nullglob
}

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

    # Clear relevant variables before sourcing to avoid pollution
    unset DRIVER_NAME CDI_KIND DRI_CONFIG MOUNT_FILES

    source "$conf_file"

    if [ -z "$DRIVER_NAME" ] || [ -z "$CDI_KIND" ]; then
        echo "WARN: Invalid config file: $conf_file (missing DRIVER_NAME or CDI_KIND), skipping"
        return
    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 no GPUs bound to this driver, delete existing config files and return
    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
    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"

    # --------------------------------------------------------
    # Step 1: Generate CDI YAML file (hostPath mode)
    # --------------------------------------------------------
    {
        echo "cdiVersion: \"${CDI_VERSION}\""
        echo "kind: \"${CDI_KIND}\""
        echo ""
        echo "devices:"
        echo "  - name: \"all\""
        echo "    containerEdits:"
        echo "      deviceNodes:"

        # Populate device nodes (hostPath mode, matches host cardx/renderDxxx index)
        for info in "${gpus[@]}"; do
            IFS=':' read -r card_name render_name _ <<< "$info"

            # Only include device nodes that exist
            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_name" ] && [ -e "/dev/dri/${render_name}" ]; then
                echo "        - hostPath: \"/dev/dri/${render_name}\""
                echo "          path: \"/dev/dri/${render_name}\""
                echo "          permissions: \"rwm\""
            fi
        done

        # Populate mount entries from MOUNT_FILES with existence checks
        echo "      mounts:"
        emit_mounts_for_driver "$DRIVER_NAME"

    } > "$cdi_file"
    chmod 644 "$cdi_file"

    # --------------------------------------------------------
    # Step 2: Generate Linglong JSON config 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

    # process all driver configs
    local count=0
    for conf in "$CONFIG_DIR"/*.conf; do
        [ -e "$conf" ] || continue
        generate_single_config "$conf"
        count=$((count + 1))
    done

    echo "INFO: Processed $count driver config(s) successfully"
}

main "$@"

