#!/bin/sh
# Clawdie operator USB live GPU setup.
# Detect display hardware and load a conservative KMS module before SDDM.
# Prefer integrated/open drivers so hybrid laptops boot the broadest path first;
# keep NVIDIA proprietary loading opt-in by requiring an installed nvidia.ko.
# Dedicated NVIDIA-target images can override that preference in rc.conf.

# PROVIDE: clawdie_live_gpu
# REQUIRE: FILESYSTEMS devfs
# BEFORE: sddm
# KEYWORD: nojail

. /etc/rc.subr

name="clawdie_live_gpu"
rcvar="clawdie_live_gpu_enable"
start_cmd="clawdie_live_gpu_start"

LOG_FILE="/var/log/clawdie-live-gpu.log"

clawdie_live_gpu_log()
{
    echo "$(date '+%H:%M:%S') $*" >> "$LOG_FILE" 2>/dev/null || true
}

clawdie_live_gpu_display_blocks()
{
    pciconf -lv 2>/dev/null | awk '
        BEGIN { RS = ""; IGNORECASE = 1 }
        /class=0x03/ { print }
    '
}

clawdie_live_gpu_nvidia_block()
{
    echo "$1" | awk '
        BEGIN { RS = ""; IGNORECASE = 1 }
        /(^|[[:space:]])vendor=0x10de([[:space:]]|$)/ { print; exit }
    '
}

clawdie_live_gpu_nvidia_device_id()
{
    _block=$(clawdie_live_gpu_nvidia_block "$1")
    [ -n "$_block" ] || return 1

    # FreeBSD pciconf chip field is chip=0x<device><vendor>. NVIDIA's vendor id
    # 0x10de sits in the LOW 16 bits, so the device id is the 4 hex digits BEFORE
    # it (e.g. chip=0x1c8c10de -> device id 1c8c). The earlier "chip=0x10de..."
    # form read the vendor half and never matched.
    echo "$_block" | grep -Eo 'chip=0x[0-9a-fA-F]{4}10de' | head -1 |
        sed 's/^chip=0x//; s/10de$//'
}

# Map an NVIDIA PCI device id (4 hex digits) to the recommended driver branch.
# Returns labels in the build's lane vocabulary {390, 470, 590} so the value can
# be compared directly with clawdie_live_gpu_nvidia_branch staged from build.sh.
# (The "590" lane installs the current nvidia-driver-580 package.)
#
# Coarse architecture heuristic by device-id range (validate against real
# hardware; refine the boundaries as needed):
#   Fermi    (       < 0x0f00) -> 390
#   Kepler   (0x0f00 .. 0x12ff) -> 470
#   Maxwell+ (      >= 0x1300) -> 590  (current: Maxwell/Pascal/Turing/Ampere/Ada/...)
# Unknown/empty -> 590, the safest default for modern unknown hardware.
clawdie_live_gpu_nvidia_branch_for_device()
{
    _device_id="$1"
    [ -n "$_device_id" ] || { echo "590"; return 0; }

    _device_num=$(printf '%d' "0x$_device_id" 2>/dev/null || echo "0")

    if [ "$_device_num" -ge 4864 ]; then
        echo "590"        # >= 0x1300  Maxwell, Pascal, Turing, Ampere, Ada, ...
    elif [ "$_device_num" -ge 3840 ]; then
        echo "470"        # 0x0f00 .. 0x12ff  Kepler
    else
        echo "390"        # < 0x0f00  Fermi
    fi
}

# Map a branch label {390,470,590} to the on-image pkg name. The "590" lane is
# the current driver, which is the nvidia-driver-580 package.
clawdie_live_gpu_nvidia_pkg_for_branch()
{
    case "$1" in
        390) echo "nvidia-driver-390" ;;
        470) echo "nvidia-driver-470" ;;
        590) echo "nvidia-driver-580" ;;
        *)   return 1 ;;
    esac
}

# nvidia-auto mode: detect the GPU, install the matching branch from the on-image
# clawdie-nvidia repo, and leave /boot/modules/nvidia.ko in place for select to
# load. Fully best-effort: any failure just logs and returns, so the caller falls
# back to the integrated/scfb path. Never blocks boot.
clawdie_live_gpu_install_nvidia()
{
    _blocks="$1"

    # Already installed (e.g. a previous boot) — nothing to do.
    if [ -e /boot/modules/nvidia.ko ]; then
        clawdie_live_gpu_log "nvidia.ko already present; skipping auto-install"
        return 0
    fi

    _device_id=$(clawdie_live_gpu_nvidia_device_id "$_blocks" || true)
    _branch=$(clawdie_live_gpu_nvidia_branch_for_device "$_device_id")
    _pkg=$(clawdie_live_gpu_nvidia_pkg_for_branch "$_branch") || {
        clawdie_live_gpu_log "WARN: no package mapping for branch ${_branch}; skipping NVIDIA auto-install"
        return 1
    }
    clawdie_live_gpu_log "nvidia-auto: device ${_device_id:-unknown} -> branch ${_branch} (${_pkg})"

    if ! command -v pkg >/dev/null 2>&1; then
        clawdie_live_gpu_log "WARN: pkg not available; cannot auto-install NVIDIA"
        return 1
    fi

    # Install only from the on-image repo so this works offline and never reaches
    # the network during boot.
    if env ASSUME_ALWAYS_YES=yes pkg install -y -r clawdie-nvidia "${_pkg}" nvidia-settings \
        >> "$LOG_FILE" 2>&1; then
        clawdie_live_gpu_log "nvidia-auto: installed ${_pkg}"
        return 0
    fi

    clawdie_live_gpu_log "WARN: nvidia-auto install of ${_pkg} failed; falling back to integrated/scfb"
    return 1
}

clawdie_live_gpu_has_pci_vendor()
{
    _blocks="$1"
    _vendor_id="$2"

    # Match the numeric PCI vendor on display-class blocks only. Do not use
    # broad text matches like "Intel" here: unrelated vendor strings or
    # subdevices can make an AMD system select i915kms and fail before SDDM.
    echo "$_blocks" | awk -v vendor="$_vendor_id" '
        BEGIN { RS = ""; IGNORECASE = 1; found = 0 }
        $0 ~ ("(^|[[:space:]])vendor=" vendor "([[:space:]]|$)") { found = 1 }
        END { exit(found ? 0 : 1) }
    '
}

clawdie_live_gpu_select_modules()
{
    _blocks="$1"
    _mode="${clawdie_live_gpu_mode:-auto}"
    _nvidia_branch="${clawdie_live_gpu_nvidia_branch:-}"

    if [ "$_mode" = "safe" ]; then
        clawdie_live_gpu_log "safe graphics mode requested; skipping KMS modules"
        return 0
    fi

    if clawdie_live_gpu_has_pci_vendor "$_blocks" "0x10de"; then
        _device_id=$(clawdie_live_gpu_nvidia_device_id "$_blocks" || true)
        _detected_branch=""
        if [ -n "$_device_id" ]; then
            _detected_branch=$(clawdie_live_gpu_nvidia_branch_for_device "$_device_id" || true)
            clawdie_live_gpu_log "detected NVIDIA device id: ${_device_id} (recommended branch: ${_detected_branch:-unknown})"
        fi

        if [ -e /boot/modules/nvidia.ko ]; then
            if [ -n "$_nvidia_branch" ] && [ -n "$_detected_branch" ] && [ "$_nvidia_branch" != "$_detected_branch" ]; then
                clawdie_live_gpu_log "WARN: staged NVIDIA branch ${_nvidia_branch} does not match detected device recommendation ${_detected_branch}; skipping proprietary load"
            elif [ "$_mode" = "nvidia" ]; then
                echo "nvidia-modeset nvidia"
                return 0
            fi
        elif [ "$_mode" = "nvidia" ]; then
            clawdie_live_gpu_log "WARN: NVIDIA mode requested but /boot/modules/nvidia.ko is missing"
        fi
    elif [ "$_mode" = "nvidia" ]; then
        clawdie_live_gpu_log "WARN: NVIDIA mode requested but no NVIDIA display device was detected"
    fi

    # Hybrid laptops often expose Intel or AMD integrated graphics plus a
    # discrete NVIDIA GPU. Prefer the integrated/open KMS path for live boot.
    if clawdie_live_gpu_has_pci_vendor "$_blocks" "0x1002"; then
        # amdgpu covers modern AMD; radeonkms is a best-effort fallback for
        # older Radeon hardware only if amdgpu does not load.
        echo "amdgpu radeonkms"
        return 0
    fi

    if clawdie_live_gpu_has_pci_vendor "$_blocks" "0x8086"; then
        echo "i915kms"
        return 0
    fi

    if clawdie_live_gpu_has_pci_vendor "$_blocks" "0x15ad"; then
        echo "vmwgfx"
        return 0
    fi

    if clawdie_live_gpu_has_pci_vendor "$_blocks" "0x10de"; then
        # NVIDIA packages are bundled in the offline repo, but the live image
        # only loads a concrete branch when one was explicitly staged.
        if [ -e /boot/modules/nvidia.ko ]; then
            echo "nvidia-modeset nvidia"
        fi
        return 0
    fi

    return 0
}

clawdie_live_gpu_load_module()
{
    _module="$1"

    if [ "$_module" = "radeonkms" ] && kldstat -n amdgpu.ko >/dev/null 2>&1; then
        clawdie_live_gpu_log "skipping radeonkms because amdgpu is already loaded"
        return 0
    fi

    if kldstat -n "${_module}.ko" >/dev/null 2>&1; then
        clawdie_live_gpu_log "module already loaded: ${_module}"
        return 0
    fi

    if kldload "$_module" >/dev/null 2>&1; then
        clawdie_live_gpu_log "loaded module: ${_module}"
        return 0
    fi

    clawdie_live_gpu_log "WARN: failed to load module: ${_module}"
    return 1
}

clawdie_live_gpu_start()
{
    _blocks=$(clawdie_live_gpu_display_blocks)
    if [ -z "$_blocks" ]; then
        clawdie_live_gpu_log "no display-class PCI device detected; using scfb/vesa fallback"
        return 0
    fi

    clawdie_live_gpu_log "detected display devices: $(echo "$_blocks" | grep -i '^vgapci' | tr '\n' ' ')"

    # Universal lane: install the matching NVIDIA branch before module selection.
    # Best-effort — on failure select_modules just falls through to integrated.
    if [ "${clawdie_live_gpu_mode:-auto}" = "nvidia-auto" ] &&
        clawdie_live_gpu_has_pci_vendor "$_blocks" "0x10de"; then
        clawdie_live_gpu_install_nvidia "$_blocks" || true
    fi

    _modules=$(clawdie_live_gpu_select_modules "$_blocks")
    if [ -z "$_modules" ]; then
        clawdie_live_gpu_log "no live KMS module selected; using scfb/vesa fallback"
        return 0
    fi

    clawdie_live_gpu_log "selected live KMS modules: ${_modules}"
    for _module in $_modules; do
        clawdie_live_gpu_load_module "$_module" || true
    done

    return 0
}

load_rc_config "$name"
: ${clawdie_live_gpu_enable:="NO"}
: ${clawdie_live_gpu_mode:="auto"}
: ${clawdie_live_gpu_nvidia_branch:=""}

run_rc_command "$1"