#!/bin/bash
set -euo pipefail

# === CONFIGURATION SECTION ==============================
# The environment overrides are intended for tests. Production uses the
# standard absolute paths below.
CONFIG_DIR="${FTG_CDI_CONFIG_DIR:-/etc/gpu-cdi/drivers.d}"
CDI_OUTPUT_DIR="${FTG_CDI_OUTPUT_DIR:-/etc/cdi}"
LINGLONG_OUTPUT_DIR="${FTG_CDI_LINGLONG_OUTPUT_DIR:-/etc/linglong/config.d}"
DRM_SYSFS_DIR="${FTG_CDI_DRM_SYSFS_DIR:-/sys/class/drm}"
DRI_DEVICE_DIR="${FTG_CDI_DRI_DEVICE_DIR:-/dev/dri}"
DRI_CONTAINER_DIR="${FTG_CDI_CONTAINER_DRI_DIR:-/dev/dri}"
CDI_VERSION="0.5.0"
DEFAULT_MOUNT_OPTS=("ro" "nosuid" "nodev" "rbind" "rprivate")
# =========================================================

detect_lib_dir() {
    local arch
    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)

ensure_output_dir() {
    local dir="$1"

    if [ -L "$dir" ]; then
        echo "ERROR: Output path must not be a symbolic link: $dir" >&2
        return 1
    fi
    mkdir -p -- "$dir"
    if [ ! -d "$dir" ] || [ ! -w "$dir" ]; then
        echo "ERROR: Output directory is not writable: $dir" >&2
        return 1
    fi
}

remove_generated_file() {
    local path="$1"

    if [ -f "$path" ] || [ -L "$path" ]; then
        echo "Removing $path"
        rm -f -- "$path"
    elif [ -e "$path" ]; then
        echo "WARN: Refusing to remove non-file path: $path" >&2
    fi
}

# Commit a generated file by rename in the same directory. A root service
# creates root-owned files; non-root test runs retain their test-user owner.
write_atomic() {
    local target="$1"
    local temp

    temp=$(mktemp "${target}.tmp.XXXXXX")
    if ! cat > "$temp"; then
        rm -f -- "$temp"
        return 1
    fi

    chmod 0644 "$temp"
    if [ "$(id -u)" -eq 0 ]; then
        chown root:root "$temp"
    fi

    # Do not replace an unchanged regular file, but always replace a symlink.
    if [ -f "$target" ] && [ ! -L "$target" ] && cmp -s "$temp" "$target"; then
        rm -f -- "$temp"
        return 0
    fi

    mv -f -- "$temp" "$target"
    chmod 0644 "$target"
}

yaml_escape() {
    local value="$1"
    value=${value//\\/\\\\}
    value=${value//\"/\\\"}
    value=${value//$'\n'/\\n}
    value=${value//$'\r'/\\r}
    printf '%s' "$value"
}

json_escape() {
    yaml_escape "$1"
}

is_indexed_array() {
    local declaration
    declaration=$(declare -p "$1" 2>/dev/null) || return 1
    [[ "$declaration" == "declare -a "* ]]
}

validate_config() {
    if [[ ! "$DRIVER_NAME" =~ ^[A-Za-z0-9_-]+$ ]]; then
        echo "ERROR: Invalid DRIVER_NAME [$DRIVER_NAME]" >&2
        return 1
    fi

    if [[ ! "$CDI_KIND" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$ ]]; then
        echo "ERROR: Invalid CDI_KIND [$CDI_KIND]" >&2
        return 1
    fi

    local path
    for path in "${MOUNT_FILES[@]}"; do
        if [[ -z "$path" || "$path" != /* || "$path" == *$'\n'* || "$path" == *$'\r'* ]]; then
            echo "ERROR: Invalid MOUNT_FILES entry [$path]" >&2
            return 1
        fi
    done

    for path in "${MOUNT_OPTS[@]}"; do
        if [[ -z "$path" || "$path" == *[[:space:]]* || "$path" == *$'\n'* || "$path" == *$'\r'* ]]; then
            echo "ERROR: Invalid MOUNT_OPTS entry [$path]" >&2
            return 1
        fi
    done

    local env_entry
    for env_entry in "${ENV_VARS[@]}"; do
        if [[ ! "$env_entry" =~ ^[A-Za-z_][A-Za-z0-9_]*=.*$ || "$env_entry" == *$'\n'* || "$env_entry" == *$'\r'* ]]; then
            echo "ERROR: Invalid ENV_VARS entry [$env_entry]" >&2
            return 1
        fi
    done
}

clean_one_config() {
    local conf="$1"
    local driver_name

    driver_name=$(sed -n \
        -e 's/^[[:space:]]*DRIVER_NAME[[:space:]]*=[[:space:]]*"\([^"]*\)".*/\1/p' \
        -e "s/^[[:space:]]*DRIVER_NAME[[:space:]]*=[[:space:]]*'\\([^']*\\)'.*/\\1/p" \
        "$conf" | head -n1)
    if [ -n "$driver_name" ] && [[ "$driver_name" =~ ^[A-Za-z0-9_-]+$ ]]; then
        remove_generated_file "${CDI_OUTPUT_DIR}/${driver_name}.yaml"
        remove_generated_file "${LINGLONG_OUTPUT_DIR}/${driver_name}.json"
    fi
}

clean_config_files() {
    if [ -d "$CONFIG_DIR" ]; then
        local conf
        for conf in "$CONFIG_DIR"/*.conf; do
            [ -e "$conf" ] || continue
            clean_one_config "$conf"
        done
    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="$DRM_SYSFS_DIR"
    [ -d "$drm_path" ] || return 0

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

        local driver_link="$card_sysfs/device/driver"
        [ -L "$driver_link" ] || continue
        local current_driver
        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
        local sorted
        IFS=$'\n' sorted=($(sort -t'|' -k3,3 -k1,1 <<<"${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:"
        local opt
        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[@]}")

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

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

count_mounts_for_driver() {
    local patterns=("${MOUNT_FILES[@]}")
    local count=0
    local pattern
    local file

    shopt -s nullglob
    for pattern in "${patterns[@]}"; do
        [ -n "$pattern" ] || continue
        for file in $pattern; do
            if [ -e "$file" ] || [ -L "$file" ]; then
                count=$((count + 1))
            fi
        done
    done
    shopt -u nullglob
    echo "$count"
}

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

    unset DRIVER_NAME CDI_KIND MOUNT_FILES MOUNT_OPTS ENV_VARS
    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=()
    elif ! is_indexed_array MOUNT_FILES; then
        echo "ERROR: MOUNT_FILES must be an indexed array in $conf_file" >&2
        return 1
    fi
    if ! declare -p MOUNT_OPTS &>/dev/null; then
        MOUNT_OPTS=("${DEFAULT_MOUNT_OPTS[@]}")
    elif ! is_indexed_array MOUNT_OPTS; then
        echo "ERROR: MOUNT_OPTS must be an indexed array in $conf_file" >&2
        return 1
    fi
    if ! declare -p ENV_VARS &>/dev/null; then
        ENV_VARS=()
    elif ! is_indexed_array ENV_VARS; then
        echo "ERROR: ENV_VARS must be an indexed array in $conf_file" >&2
        return 1
    fi
    validate_config

    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"
        remove_generated_file "$cdi_file"
        remove_generated_file "$linglong_file"
        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 "---"
        echo "cdiVersion: \"${CDI_VERSION}\""
        echo "kind: \"${CDI_KIND}\""
        echo ""
        echo "devices:"
        echo "  - name: \"all\""
        echo "    containerEdits:"

        if [ ${#ENV_VARS[@]} -gt 0 ]; then
            echo "      env:"
            local env_entry
            for env_entry in "${ENV_VARS[@]}"; do
                echo "        - \"$(yaml_escape "$env_entry")\""
            done
        fi

        local device_nodes=()
        local node_host node_container
        local info
        for info in "${gpus[@]}"; do
            local card_name render_names sort_key
            IFS='|' read -r card_name render_names sort_key <<< "$info"
            if [ -e "${DRI_DEVICE_DIR}/${card_name}" ]; then
                device_nodes+=("${DRI_DEVICE_DIR}/${card_name}|${DRI_CONTAINER_DIR}/${card_name}")
            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 "${DRI_DEVICE_DIR}/${render_name}" ]; then
                        device_nodes+=("${DRI_DEVICE_DIR}/${render_name}|${DRI_CONTAINER_DIR}/${render_name}")
                    fi
                done
            fi
        done

        echo "      deviceNodes:"
        if [ ${#device_nodes[@]} -eq 0 ]; then
            echo "        []"
        else
            for node in "${device_nodes[@]}"; do
                IFS='|' read -r node_host node_container <<< "$node"
                echo "        - hostPath: \"${node_host}\""
                echo "          path: \"${node_container}\""
                echo "          permissions: \"rwm\""
            done
        fi

        local mount_count
        mount_count=$(count_mounts_for_driver)
        if [ "$mount_count" -eq 0 ]; then
            echo "      mounts: []"
        else
            echo "      mounts:"
            emit_mounts_for_driver "$DRIVER_NAME"
        fi
    } | write_atomic "$cdi_file"

    {
        echo "{"
        echo "    \"devices\": ["
        echo "        \"$(json_escape "${CDI_KIND}=all")\""
        echo "    ]"
        echo "}"
    } | write_atomic "$linglong_file"
}

main() {
    if [ "${1:-}" = "--clean" ]; then
        echo "INFO: Cleaning all generated CDI configuration files..."
        clean_config_files
        return 0
    fi

    if [ ! -d "$CONFIG_DIR" ]; then
        echo "INFO: Driver config directory $CONFIG_DIR not found, exiting"
        return 0
    fi

    echo "INFO: Cleaning previous CDI configuration files..."
    clean_config_files
    ensure_output_dir "$CDI_OUTPUT_DIR"
    ensure_output_dir "$LINGLONG_OUTPUT_DIR"

    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 "$@"
