#!/bin/bash
# mthreads-ctk — Moore Threads container toolkit CLI (CDI 0.5.0 spec generator).
#
# Scans /sys/class/drm for GPUs bound to the kernel driver declared in each
# /etc/mthreads-ctk/drivers.d/<driver>.conf and writes:
#   /etc/cdi/<driver>.yaml              CDI spec (read by podman/docker/containerd/linyaps)
#   /etc/linglong/config.d/<driver>.json linyaps device enablement (only when linyaps is present)
#
# Usage: mthreads-ctk generate      regenerate specs for all configured drivers (default)
#        mthreads-ctk clean         remove generated specs
#        mthreads-ctk --help
# Environment overrides (testing):
#   MTHREADS_CTK_CONFIG_DIR, MTHREADS_CTK_OUTPUT_DIR, MTHREADS_CTK_LINGLONG_DIR, MTHREADS_CTK_SYSFS_DRM
#   MTHREADS_CTK_LINGLONG=auto|always|never   (default auto: detect ll-cli or /etc/linglong)
set -euo pipefail

CONFIG_DIR="${MTHREADS_CTK_CONFIG_DIR:-/etc/mthreads-ctk/drivers.d}"
CDI_OUTPUT_DIR="${MTHREADS_CTK_OUTPUT_DIR:-/etc/cdi}"
LINGLONG_OUTPUT_DIR="${MTHREADS_CTK_LINGLONG_DIR:-/etc/linglong/config.d}"
SYSFS_DRM="${MTHREADS_CTK_SYSFS_DRM:-/sys/class/drm}"
LINGLONG_MODE="${MTHREADS_CTK_LINGLONG:-auto}"
CDI_VERSION="0.5.0"
MOUNT_OPTS=("ro" "nosuid" "nodev" "rbind" "rprivate")

log() { echo "INFO: $*"; }
warn() { echo "WARN: $*" >&2; }
err() { echo "ERROR: $*" >&2; }

detect_lib_dir() {
    local arch
    arch=$(uname -m)
    local triplet=""
    case "$arch" in
        x86_64) triplet="x86_64-linux-gnu" ;;
        aarch64) triplet="aarch64-linux-gnu" ;;
        loongarch64) triplet="loongarch64-linux-gnu" ;;
    esac
    if [ -n "$triplet" ] && [ -d "/usr/lib/$triplet" ]; then
        echo "/usr/lib/$triplet"
    elif [ -d /usr/lib64 ]; then
        echo "/usr/lib64"
    else
        echo "/usr/lib"
    fi
}
# shellcheck disable=SC2034  # consumed by the sourced drivers.d/*.conf (${LIB_DIR}/...)
LIB_DIR=$(detect_lib_dir)

# linyaps enablement JSON is only meaningful where linyaps is installed.
linglong_wanted() {
    case "$LINGLONG_MODE" in
        always) return 0 ;;
        never) return 1 ;;
        *) [ -x /usr/bin/ll-cli ] || [ -d /etc/linglong ] ;;
    esac
}

# Atomic replace: write to a same-directory temp file, fix ownership/mode, rename.
atomic_install() {
    local src="$1" dst="$2"
    chmod 0644 "$src"
    if [ "$(id -u)" -eq 0 ]; then
        chown root:root "$src"
    fi
    mv -f "$src" "$dst"
}

remove_if_present() {
    local f="$1"
    if [ -e "$f" ] || [ -L "$f" ]; then
        log "Removing $f"
        rm -f "$f"
    fi
}

# Reject names that could escape the output directories or break YAML/JSON.
valid_driver_name() { [[ "$1" =~ ^[A-Za-z0-9_-]+$ ]]; }
valid_cdi_kind() { [[ "$1" =~ ^[a-z0-9]([a-z0-9.-]*[a-z0-9])?/[A-Za-z0-9_-]+$ ]]; }

clean_all() {
    log "Cleaning all generated CDI configuration files..."
    local conf driver_name
    for conf in "$CONFIG_DIR"/*.conf; do
        [ -e "$conf" ] || continue
        driver_name=$(sed -n 's/^[[:space:]]*DRIVER_NAME=["'"'"']\{0,1\}\([A-Za-z0-9_-]*\)["'"'"']\{0,1\}.*/\1/p' "$conf" | head -n1)
        [ -n "$driver_name" ] || continue
        remove_if_present "${CDI_OUTPUT_DIR}/${driver_name}.yaml"
        remove_if_present "${LINGLONG_OUTPUT_DIR}/${driver_name}.json"
    done
}

# Output lines: card_name|render_names|sort_key ; sorted by sort_key (PCI address, sysfs fallback)
scan_gpus() {
    local driver_name=$1
    local -n gpu_arr=$2
    gpu_arr=()
    [ -d "$SYSFS_DRM" ] || return 0

    local card_sysfs card_name driver_link current_driver render_names dev sort_key resolved
    for card_sysfs in "$SYSFS_DRM"/card[0-9]*; do
        [ -e "$card_sysfs" ] || continue
        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

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

        sort_key="$card_sysfs"
        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=()
        mapfile -t sorted < <(printf '%s\n' "${gpu_arr[@]}" | sort -t'|' -k3,3)
        gpu_arr=("${sorted[@]}")
    fi
}

emit_mount() {
    local path="$1"
    echo "        - hostPath: \"${path}\""
    echo "          containerPath: \"${path}\""
    echo "          options:"
    local opt
    for opt in "${MOUNT_OPTS[@]}"; do
        echo "            - ${opt}"
    done
}

# Expand MOUNT_FILES globs; emit each existing path once.
emit_mounts_for_driver() {
    local driver_name="$1"
    local patterns=(${MOUNT_FILES[@]+"${MOUNT_FILES[@]}"})
    if [ ${#patterns[@]} -eq 0 ]; then
        warn "MOUNT_FILES is empty for driver [$driver_name], no mounts will be generated"
        return 0
    fi
    local -A seen=()
    local pattern file matched
    shopt -s nullglob
    for pattern in "${patterns[@]}"; do
        [ -n "$pattern" ] || continue
        matched=0
        for file in $pattern; do
            matched=1
            # A dangling symlink must not be emitted: the runtime bind-mounts the resolved
            # target and fails the whole container start when it does not exist.
            if [ ! -e "$file" ]; then
                [ -L "$file" ] && warn "Skipping dangling symlink: $file"
                continue
            fi
            if [ -z "${seen[$file]:-}" ]; then
                seen[$file]=1
                emit_mount "$file"
            fi
        done
        [ "$matched" -eq 1 ] || warn "Skipping mount for non-existent path: $pattern"
    done
    shopt -u nullglob
}

generate_single_config() {
    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
        err "Invalid config file: $conf_file (missing DRIVER_NAME or CDI_KIND)"
        return 1
    fi
    if ! valid_driver_name "$DRIVER_NAME"; then
        err "Invalid DRIVER_NAME in $conf_file: [$DRIVER_NAME] (allowed: A-Z a-z 0-9 _ -)"
        return 1
    fi
    if ! valid_cdi_kind "$CDI_KIND"; then
        err "Invalid CDI_KIND in $conf_file: [$CDI_KIND] (expected <vendor.domain>/<class>)"
        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
        log "No GPU device bound to driver [$DRIVER_NAME], removing old config if any"
        remove_if_present "$cdi_file"
        remove_if_present "$linglong_file"
        return 0
    fi

    log "Driver [$DRIVER_NAME] - found $gpu_count GPU(s)"
    log "  Generating CDI config: $cdi_file"

    mkdir -p "$CDI_OUTPUT_DIR"
    local tmp
    tmp=$(mktemp "${CDI_OUTPUT_DIR}/.${DRIVER_NAME}.yaml.XXXXXX")
    {
        echo "cdiVersion: \"${CDI_VERSION}\""
        echo "kind: \"${CDI_KIND}\""
        echo ""
        echo "devices:"
        echo "  - name: \"all\""
        echo "    containerEdits:"
        echo "      deviceNodes:"
        local info card_name render_names render_name render_list
        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
                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"
    } >"$tmp"
    atomic_install "$tmp" "$cdi_file"

    if linglong_wanted; then
        log "  Generating Linglong config: $linglong_file"
        mkdir -p "$LINGLONG_OUTPUT_DIR"
        tmp=$(mktemp "${LINGLONG_OUTPUT_DIR}/.${DRIVER_NAME}.json.XXXXXX")
        {
            echo "{"
            echo "    \"devices\": ["
            echo "        \"${CDI_KIND}=all\""
            echo "    ]"
            echo "}"
        } >"$tmp"
        atomic_install "$tmp" "$linglong_file"
    else
        log "  Linglong not detected, skipping $linglong_file"
        remove_if_present "$linglong_file"
    fi
}

usage() {
    echo "Usage: mthreads-ctk generate | clean | --help"
    echo "  generate  regenerate CDI specs for all drivers declared in $CONFIG_DIR (default)"
    echo "  clean     remove generated CDI specs and linyaps enablement files"
}

main() {
    local cmd="${1:-generate}"
    case "$cmd" in
        clean|--clean)
            clean_all
            exit 0
            ;;
        generate) ;;
        -h|--help|help)
            usage
            exit 0
            ;;
        *)
            err "unknown command: $cmd"
            usage >&2
            exit 64
            ;;
    esac
    if [ ! -d "$CONFIG_DIR" ]; then
        log "Driver config directory $CONFIG_DIR not found, exiting"
        exit 0
    fi

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

    log "Processed $count driver config(s) successfully"
    if [ "$failures" -gt 0 ]; then
        err "$failures driver config(s) failed"
        return 1
    fi
}

main "$@"
