#!/bin/sh
# Clawdie Shell — NVIDIA GPU Detection Module
# Purpose: Detect NVIDIA GPU and select appropriate driver version
# POSIX-compliant (no bash-isms)

set -eu

# Configuration (can be overridden for testing)
RC_CONF="${RC_CONF:-/etc/rc.conf}"
LOG_FILE="${LOG_FILE:-/var/log/clawdie-firstboot.log}"
PROGRESS_FILE="${PROGRESS_FILE:-/var/log/clawdie-firstboot.progress}"

# ============================================================================
# MAIN ENTRY POINT
# ============================================================================

clawdie_shell_nvidia_detect() {
  # Main orchestrator: detect NVIDIA GPU and select driver version
  # Only runs if DETECTED_GPU=nvidia (called by firstboot.sh)
  # Sets NVIDIA_DRIVER to 590, 470, or 390

  if [ "${DETECTED_GPU:-}" != "nvidia" ]; then
    log_msg "[nvidia] Not an NVIDIA GPU — skipping"
    return 0
  fi

  log_msg "[nvidia] Detecting NVIDIA GPU type"

  # Query pciconf for NVIDIA device ID
  local device_id
  device_id=$(clawdie_shell_nvidia_get_device_id) || {
    log_msg "[nvidia] ERROR: Could not detect NVIDIA device ID"
    NVIDIA_DRIVER="590"
    echo "[NVIDIA] Default 590" >> "$PROGRESS_FILE"
    return 1
  }

  log_msg "[nvidia] Device ID: $device_id"

  # Match device ID to driver version
  NVIDIA_DRIVER=$(clawdie_shell_nvidia_select_driver "$device_id")

  log_msg "[nvidia] Selected driver version: nvidia-driver-$NVIDIA_DRIVER"

  # Update rc.conf to include correct NVIDIA driver version
  clawdie_shell_nvidia_write_rcconf "$NVIDIA_DRIVER"

  # Install the NVIDIA driver package
  clawdie_shell_nvidia_install "$NVIDIA_DRIVER"

  echo "[NVIDIA] $NVIDIA_DRIVER" >> "$PROGRESS_FILE"
  log_msg "[nvidia] NVIDIA detection complete"
}

# ============================================================================
# NVIDIA DEVICE DETECTION
# ============================================================================

clawdie_shell_nvidia_get_device_id() {
  # Query pciconf for NVIDIA GPU device ID
  # Returns: device ID in hex format (e.g., "2388")

  local pciconf_out
  local device_id

  pciconf_out=$(pciconf -lv 2>/dev/null | grep -i "class=0x030000.*NVIDIA" | head -1 || true)

  if [ -z "$pciconf_out" ]; then
    return 1
  fi

  # Extract device ID from chip field (format: chip=0xVVDDxx)
  # where VV = vendor (10de for NVIDIA), DD = device ID, xx = revision
  device_id=$(echo "$pciconf_out" | grep -o "chip=0x10de[a-f0-9]*" | cut -c11-14 || true)

  if [ -z "$device_id" ]; then
    return 1
  fi

  echo "$device_id"
}

# ============================================================================
# DRIVER VERSION SELECTION
# ============================================================================

clawdie_shell_nvidia_select_driver() {
  # Map NVIDIA device ID to supported driver version
  # Input: device ID (hex, 4 chars)
  # Output: driver version (590, 470, or 390)
  #
  # Rough mapping (simplified for common GPUs):
  # - Maxwell (GTX 750, 750 Ti, GTX 950-980 Ti) → 470 (0x1340-0x2186)
  # - Kepler (GTX 650-780 Ti) → 390 (0x1180-0x139D)
  # - Turing/Ampere (RTX 2060+, RTX 30xx) → 590 (0x2060+)
  # - Ada (RTX 40xx) → 590 (0x2700+)

  local device_id="${1:-0000}"
  local device_num

  # Convert hex device ID to decimal for easier comparison
  device_num=$(printf "%d" "0x$device_id" 2>/dev/null || echo "0")

  # Turing and newer (RTX 20xx, 30xx, 40xx, etc.) — driver 590
  # RTX 2060+ starts around 0x2060 (8288), RTX 30xx around 0x2200 (8704), RTX 40xx around 0x2700 (9984)
  if [ "$device_num" -ge 8288 ]; then
    echo "590"
    return 0
  fi

  # Maxwell (GTX 750-980 Ti) — driver 470
  # Range roughly 0x1340 (4928) to 0x2186 (8582)
  if [ "$device_num" -ge 4928 ] && [ "$device_num" -le 8582 ]; then
    echo "470"
    return 0
  fi

  # Kepler and older (GTX 650-780 Ti) — driver 390
  # Range roughly 0x1180 (4480) to 0x139F (5023)
  if [ "$device_num" -ge 4480 ] && [ "$device_num" -le 5100 ]; then
    echo "390"
    return 0
  fi

  # Default: assume modern GPU → driver 590
  echo "590"
}

# ============================================================================
# RC.CONF WRITING
# ============================================================================

clawdie_shell_nvidia_write_rcconf() {
  # Write NVIDIA driver version to rc.conf
  # Input: driver version (590, 470, or 390)

  local nvidia_version="${1:-590}"
  local pkg_name="nvidia-driver-$nvidia_version"

  if [ ! -f "$RC_CONF" ]; then
    log_msg "[nvidia] Creating $RC_CONF"
    touch "$RC_CONF"
  fi

  # Remove existing NVIDIA references from kld_list
  if grep -q "^kld_list=" "$RC_CONF" 2>/dev/null; then
    # Remove nvidia and nvidia-modeset from existing kld_list
    local kld=$(grep "^kld_list=" "$RC_CONF" | cut -d= -f2 | tr -d '"')
    kld=$(echo "$kld" | sed 's/nvidia-modeset //g; s/nvidia //g; s/  / /g' | xargs)

    sed -i.bak '/^kld_list=/d' "$RC_CONF"
    rm -f "$RC_CONF.bak"

    # Re-add kld_list with NVIDIA driver
    if [ -n "$kld" ]; then
      echo "kld_list=\"nvidia-modeset nvidia $kld\"" >> "$RC_CONF"
    else
      echo "kld_list=\"nvidia-modeset nvidia\"" >> "$RC_CONF"
    fi
  else
    # No existing kld_list
    echo "kld_list=\"nvidia-modeset nvidia\"" >> "$RC_CONF"
  fi

  log_msg "[nvidia] Configured rc.conf for $pkg_name"
}

# ============================================================================
# PACKAGE INSTALLATION
# ============================================================================

clawdie_shell_nvidia_install() {
  # Install NVIDIA driver package
  # Input: driver version (590, 470, or 390)

  local nvidia_version="${1:-590}"
  local pkg_name="nvidia-driver-$nvidia_version"

  # Check if already installed
  if pkg info "$pkg_name" >/dev/null 2>&1; then
    log_msg "[nvidia] $pkg_name already installed"
    return 0
  fi

  log_msg "[nvidia] Installing $pkg_name..."

  # Try to install (works offline if packages on USB, online if network available)
  if pkg install -y "$pkg_name" >/dev/null 2>&1; then
    log_msg "[nvidia] Successfully installed $pkg_name"
    return 0
  else
    log_msg "[nvidia] WARNING: Failed to install $pkg_name"
    log_msg "[nvidia] System may need network access or USB packages"
    return 1
  fi
}

# ============================================================================
# VALIDATION
# ============================================================================

clawdie_shell_nvidia_validate() {
  # Verify NVIDIA configuration in rc.conf

  if [ ! -f "$RC_CONF" ]; then
    log_msg "[nvidia] ERROR: $RC_CONF not found"
    return 1
  fi

  if ! grep -q "nvidia" "$RC_CONF" 2>/dev/null; then
    log_msg "[nvidia] ERROR: NVIDIA not configured in rc.conf"
    return 1
  fi

  return 0
}

# ============================================================================
# LOGGING HELPER
# ============================================================================

log_msg() {
  echo "$(date '+%H:%M:%S') $1" | tee -a "$LOG_FILE" 2>/dev/null || true
}

# Only run if sourced directly (not during test)
if [ "${SHELL_NVIDIA_TEST:-0}" -eq 0 ]; then
  clawdie_shell_nvidia_detect
fi