#!/bin/bash
#
# Copy missing vendor-signed Qualcomm firmware from a Windows driver store.
# Required filenames come from the running device tree.
#
# Modes
#   --stage DIR       Live ISO, before the disk is touched: scan the Windows
#                     partitions and save the missing files under DIR (the
#                     installer copies DIR into the target). Missing firmware
#                     is reported without aborting staging. A rerun adds the
#                     files earlier runs did not find.
#   --install         Installed system or the installer's chroot: take each
#                     file from the stage (--stage-dir, default
#                     /var/lib/omarchy/firmware-stage), else from a Windows
#                     partition still on an internal disk, write them under
#                     /usr/lib/firmware/updates, record a manifest, add the
#                     zap shader to the initramfs, and rebuild the boot image
#                     (unless --no-rebuild). This is the default mode.
#   --list-missing    Print the device-tree firmware names not present under
#                     /usr/lib/firmware and not listed for this board (below),
#                     one per line. Needs no root.
#   -d DIR            Use DIR as the driver store (a Windows/System32/
#                     DriverStore/FileRepository directory, or any directory
#                     tree holding the files) instead of the Windows
#                     partitions. --install then ignores the stage.
#
# Files land in /usr/lib/firmware/updates to avoid package ownership conflicts.
#
# A board package lists named firmware that must not come from Windows on its
# board, one name per line, in /usr/share/qcom-firmware-extract/provided.d/
# <device-tree compatible>.list: files the package handles itself, and files
# the machine's own boot firmware loads before Linux starts. Names listed
# for one of this machine's compatibles are never searched for, installed or
# reported missing. When nothing else is missing, the disks are not touched.
# A copy that an earlier run installed before the name was listed is left in
# place; the next --install stops recording it and prints where it is. Since
# --list-missing leaves them out, a board must not list firmware that a
# hardware check relies on, such as the ADSP image Omarchy checks before it
# enables the DSP driver, unless its package installs that file before the
# check runs.
#
# Derived from Canonical's qcom-firmware-extract (Tobias Heider, GPL-2+):
# https://salsa.debian.org/debian/qcom-firmware-extract
# SPDX-License-Identifier: GPL-2.0-or-later

set -euo pipefail

DT_ROOT=${QCOM_FW_DT_ROOT:-/sys/firmware/devicetree/base}
FW_ROOT=${QCOM_FW_FIRMWARE_ROOT:-/usr/lib/firmware}
KERNEL_CONFIG=${QCOM_FW_KERNEL_CONFIG:-/proc/config.gz}
PROVIDED_DIR=${QCOM_FW_PROVIDED_DIR:-/usr/share/qcom-firmware-extract/provided.d}
# Destination prefix for the installed-system and mount paths (tests only).
ROOT=${QCOM_FW_ROOT:-}
# The effective user ID the root checks see (tests only).
CHECKED_EUID=${QCOM_FW_TEST_EUID:-$EUID}
STATE_DIR="$ROOT/var/lib/omarchy/qcom-firmware"
MANIFEST="$STATE_DIR/manifest"
INITCPIO_CONF="$ROOT/etc/mkinitcpio.conf.d/qcom-firmware.conf"
DEFAULT_STAGE_DIR="$ROOT/var/lib/omarchy/firmware-stage"
DRIVERSTORE_PATH="Windows/System32/DriverStore/FileRepository"

mode=install
stage_dir=""
driverstore=""
rebuild=yes

usage() {
  sed -n '2,/^# SPDX/{/^# SPDX/d;s/^# \{0,1\}//p}' "$0"
}

while (($#)); do
  case $1 in
    --stage) mode=stage; stage_dir=${2:?--stage needs a directory}; shift ;;
    --stage-dir) stage_dir=${2:?--stage-dir needs a directory}; shift ;;
    --install) mode=install ;;
    --list-missing) mode=list ;;
    --no-rebuild) rebuild=no ;;
    -d) driverstore=${2:?-d needs a directory}; shift ;;
    -h | --help) usage; exit 0 ;;
    *) echo "qcom-firmware-extract: unknown option: $1" >&2; usage >&2; exit 2 ;;
  esac
  shift
done

log() { printf 'qcom-firmware-extract: %s\n' "$*"; }
warn() { printf 'qcom-firmware-extract: %s\n' "$*" >&2; }

# ---------------------------------------------------------------------------
# 1. What does this machine ask for, and what is already there?
# ---------------------------------------------------------------------------

# Prints "<node-name> <firmware path>" for every firmware-name entry in the
# device tree. A property can hold several NUL-separated strings.
dt_firmware_names() {
  local prop node name
  [[ -d $DT_ROOT ]] || return 0
  find "$DT_ROOT" -name firmware-name -type f -print0 2>/dev/null | sort -z |
    while IFS= read -r -d '' prop; do
      node=${prop%/firmware-name}
      node=${node#"$DT_ROOT"/}
      tr '\0' '\n' <"$prop" | while IFS= read -r name; do
        [[ -n $name ]] && printf '%s %s\n' "$node" "$name"
      done
    done
}

# The kernel tries plain firmware in every directory, then zstd and xz if built
# with them. gzip is not an automatic firmware-loader fallback.
firmware_suffixes=('' .zst .xz)
if kernel_config=$(gzip -dc -- "$KERNEL_CONFIG" 2>/dev/null); then
  firmware_suffixes=('')
  if grep -qx 'CONFIG_FW_LOADER_COMPRESS_ZSTD=y' <<<"$kernel_config"; then
    firmware_suffixes+=(.zst)
  fi
  if grep -qx 'CONFIG_FW_LOADER_COMPRESS_XZ=y' <<<"$kernel_config"; then
    firmware_suffixes+=(.xz)
  fi
fi

firmware_file() {
  local name=$1 dir ext
  for ext in "${firmware_suffixes[@]}"; do
    for dir in "$FW_ROOT/updates" "$FW_ROOT"; do
      if [[ -f "$dir/$name$ext" ]]; then
        printf '%s\n' "$dir/$name$ext"
        return 0
      fi
    done
  done
  return 1
}

# Names that board packages list for this machine as not from Windows.
declare -A provided=()
load_provided_names() {
  local compatible list line
  [[ -r $DT_ROOT/compatible && -d $PROVIDED_DIR ]] || return 0
  while IFS= read -r compatible; do
    [[ -n $compatible && $compatible != */* ]] || continue
    list="$PROVIDED_DIR/$compatible.list"
    [[ -f $list ]] || continue
    if [[ ! -r $list ]]; then
      warn "ignoring $list: it cannot be read"
      continue
    fi
    while IFS= read -r line || [[ -n $line ]]; do
      # A comment starts at a '#' that begins the line or follows a space.
      line=" $line"
      line=${line%%[[:space:]]#*}
      line=${line#"${line%%[![:space:]]*}"}
      line=${line%"${line##*[![:space:]]}"}
      [[ -n $line ]] || continue
      if [[ $line == *[[:space:]]* ]]; then
        warn "ignoring '$line' in $list: a firmware name has no spaces"
        continue
      fi
      provided[$line]=$compatible
    done <"$list"
  done < <(tr '\0' '\n' <"$DT_ROOT/compatible")
  return 0
}
load_provided_names

declare -a missing_nodes=() missing_names=() provided_names=()
while read -r node name; do
  if [[ -n ${provided[$name]-} ]]; then
    provided_names+=("$name")
    continue
  fi
  firmware_file "$name" >/dev/null && continue
  missing_nodes+=("$node")
  missing_names+=("$name")
done < <(dt_firmware_names)

if [[ $mode == list ]]; then
  ((${#missing_names[@]})) && printf '%s\n' "${missing_names[@]}"
  exit 0
fi

if [[ ! -d $DT_ROOT ]]; then
  log "no device tree on this machine, nothing to do"
  exit 0
fi
for name in "${provided_names[@]}"; do
  log "skipping $name: listed in $PROVIDED_DIR/${provided[$name]}.list"
done
if [[ $mode != install ]] && ((${#missing_names[@]} == 0)); then
  log "every firmware file the device tree names is already installed${provided_names[*]:+ or listed for this board}"
  exit 0
fi

# Root is needed to write under /usr/lib/firmware (and, below, to mount
# partitions); staging a -d driver store into a user directory needs neither.
if [[ $CHECKED_EUID != 0 && $mode == install ]]; then
  warn "must run as root"
  exit 1
fi

# ---------------------------------------------------------------------------
# 2. Where can they come from?
# ---------------------------------------------------------------------------

declare -a sources=()   # driver stores to search together (-d or Windows)
declare -a mounts=()    # temporary mounts to undo
stage_source=""         # the installer's stage, used before the driver stores
declare -A manifest=()  # the manifest this run updates, lines by firmware name
declare -A staged=()    # names the stage manifest lists
cleanup() {
  local m
  for m in "${mounts[@]-}"; do
    [[ -n $m ]] || continue
    umount "$m" 2>/dev/null || true
    rmdir "$m" 2>/dev/null || true
  done
}
trap cleanup EXIT

# Mount Windows driver stores read-only, preferring internal disks. With
# "internal-only", removable and external disks are left alone.
scan_windows_partitions() {
  local scope=${1:-all} dev removable hotplug fstype tran mp bitlocker="" skipped=""
  if [[ $CHECKED_EUID != 0 ]]; then
    warn "must run as root to read the Windows partitions"
    exit 1
  fi
  local -a internal=() external=()
  # read merges empty columns, so the always-set RM and HOTPLUG come first
  # and an empty TRAN comes last.
  while read -r dev removable hotplug fstype tran; do
    [[ $fstype == ntfs || $fstype == BitLocker ]] || continue
    if [[ $removable == 0 && $hotplug == 0 && $tran != usb ]]; then
      if [[ $fstype == ntfs ]]; then internal+=("$dev"); else bitlocker+=" $dev"; fi
    elif [[ $scope == internal-only ]]; then
      skipped+=" $dev"
    elif [[ $fstype == ntfs ]]; then
      external+=("$dev")
    else
      bitlocker+=" $dev"
    fi
  done < <(lsblk -rpno PATH,RM,HOTPLUG,FSTYPE,TRAN 2>/dev/null)
  [[ -z $skipped ]] || log "not reading removable or external disk partition(s):$skipped; use -d for a driver store on one"

  for dev in "${internal[@]}" "${external[@]}"; do
    mp=$(mktemp -d "$ROOT/run/qcom-firmware-extract.XXXXXX")
    if mount -t ntfs3 -o ro,nosuid,nodev,noexec "$dev" "$mp" 2>/dev/null; then
      mounts+=("$mp")
      if [[ -d "$mp/$DRIVERSTORE_PATH" ]]; then
        log "Windows driver store found on $dev"
        sources+=("$mp/$DRIVERSTORE_PATH")
      fi
    else
      warn "could not mount $dev read-only (ntfs3); skipping it"
      rmdir "$mp" 2>/dev/null || true
    fi
  done
  [[ -n $bitlocker ]] && warn "BitLocker volume(s) skipped:$bitlocker; turn BitLocker off in Windows to let the firmware be read"
  return 0
}

firmware_sha256() {
  local file=$1
  case "$file" in
    *.zst) zstd -qdc -- "$file" ;;
    *.xz) xz -dc -- "$file" ;;
    *) cat -- "$file" ;;
  esac | sha256sum | cut -d' ' -f1
}

# Match ambiguous DTBs to an installed sibling firmware blob from the same node.
matching_companion_candidate() {
  local node=$1
  shift
  local -a candidates=("$@")
  local dt_node dt_name installed installed_hash candidate sibling matched compatible
  local -a matches=()

  for candidate in "${candidates[@]}"; do
    matched=0
    compatible=1
    while read -r dt_node dt_name; do
      [[ $dt_node == "$node" ]] || continue
      installed=$(firmware_file "$dt_name") || continue
      installed_hash=$(firmware_sha256 "$installed") || return 1
      sibling=$(find "${candidate%/*}" -maxdepth 1 -type f \
        -iname "$(basename "$dt_name")" -print -quit 2>/dev/null)
      [[ -n $sibling ]] || continue
      if [[ $(sha256sum "$sibling" | cut -d' ' -f1) == "$installed_hash" ]]; then
        matched=1
      else
        compatible=0
        break
      fi
    done < <(dt_firmware_names)
    if ((matched && compatible)); then
      matches+=("$candidate")
    fi
  done

  unique_candidate "${matches[@]}"
}

# Duplicate copies are safe only when their contents agree.
unique_candidate() {
  (($#)) || return 1
  local first=$1 candidate
  shift
  for candidate in "$@"; do
    cmp -s "$first" "$candidate" || return 1
  done
  printf '%s\n' "$first"
}

# Windows can ship DSP device-tree images under another name: a device tree
# may ask for adsp_dtb.mbn where Windows has adsp_dtbs.elf (e.g. Surface Pro 11).
windows_firmware_names() {
  printf '%s\n' "$1"
  if [[ $1 == *_dtb.mbn ]]; then
    printf '%s\n' "${1%_dtb.mbn}_dtbs.elf"
  fi
}

# Require a companion match when several firmware variants exist. Returns 1
# when the sources have no candidate and 2 when it refuses ambiguous ones.
find_in_sources() {
  local base=$1 node=$2 candidate name
  local -a candidates=()
  while IFS= read -r name; do
    while IFS= read -r -d '' candidate; do
      candidates+=("$candidate")
    done < <(find "${sources[@]}" -type f -iname "$name" -print0 2>/dev/null)
    ((${#candidates[@]})) && break
  done < <(windows_firmware_names "$base")

  ((${#candidates[@]})) || return 1
  if unique_candidate "${candidates[@]}"; then
    return 0
  fi

  if matching_companion_candidate "$node" "${candidates[@]}"; then
    return 0
  fi

  warn "refusing ambiguous $base: no unique variant matches its installed companion firmware"
  return 2
}

# A manifest has one line per firmware name: the name, the SHA-256 of the file
# and where it came from. Loads FILE as the manifest this run updates.
read_manifest() {
  local line name
  manifest=()
  while IFS= read -r line || [[ -n $line ]]; do
    read -r name _ <<<"$line"
    if [[ -n $name ]]; then
      manifest[$name]=$line
    fi
  done <"$1"
}

# Sets NAME's manifest line to NAME followed by FIELDS; without FIELDS, removes
# it.
set_manifest_line() {
  local name=$1
  shift
  if (($#)); then
    manifest[$name]="$name $*"
  else
    unset 'manifest[$name]'
  fi
}

write_manifest() {
  if ((${#manifest[@]})); then
    printf '%s\n' "${manifest[@]}" | sort >"$1.new"
  else
    : >"$1.new"
  fi
  mv "$1.new" "$1"
}

# A file in DIR counts as staged when the stage manifest lists it. In --stage
# mode, that manifest is also the one the run updates.
use_stage() {
  local name
  stage_source=$1
  read_manifest "$1/manifest"
  for name in "${!manifest[@]}"; do
    staged[$name]=1
  done
}

is_staged() {
  [[ -n ${staged[$1]-} && -f $stage_source/$1 ]]
}

# True if the stage lacks any of the missing files.
stage_incomplete() {
  local name
  for name in "${missing_names[@]}"; do
    is_staged "$name" || return 0
  done
  return 1
}

# The stage comes first; the driver stores supply what it lacks. A staged file
# listed under another device-tree path with the same file name (a later device
# tree may move it) is the last resort, since after a full-disk install the
# stage is the only copy.
locate_firmware() {
  local name=$1 node=$2 src status=0
  if is_staged "$name"; then
    printf '%s\n' "$stage_source/$name"
    return 0
  fi
  if ((${#sources[@]})); then
    src=$(find_in_sources "$(basename "$name")" "$node") || status=$?
    if ((status == 0)); then
      printf '%s\n' "$src"
      return 0
    fi
    # The stage does not override a refusal of ambiguous variants.
    ((status != 2)) || return 1
  fi
  [[ -n $stage_source ]] || return 1
  local -a sources=("$stage_source")
  src=$(find_in_sources "$(basename "$name")" "$node") || return 1
  is_staged "${src#"$stage_source"/}" || return 1
  printf '%s\n' "$src"
}

# ---------------------------------------------------------------------------
# 3. Stage: live session, before the disk is written
# ---------------------------------------------------------------------------

if [[ $mode == stage ]]; then
  # Files staged by an earlier run are kept; a rerun, for example with -d,
  # looks only for the rest, so an empty or partial stage blocks nothing.
  if [[ -f "$stage_dir/manifest" ]]; then
    use_stage "$stage_dir"
  fi
  if ! stage_incomplete; then
    log "already staged in $stage_dir"
    exit 0
  fi
  if [[ -n $driverstore ]]; then
    sources=("$driverstore")
  elif [[ -n $stage_source && $CHECKED_EUID != 0 ]]; then
    # Without root the partitions cannot be read; keep the stage as it is.
    log "already staged in $stage_dir; run as root to read the Windows partitions for the rest"
    exit 0
  else
    scan_windows_partitions
  fi
  mkdir -p "$stage_dir"
  copied=0
  kept=0
  for i in "${!missing_names[@]}"; do
    node=${missing_nodes[$i]}
    name=${missing_names[$i]}
    if is_staged "$name"; then
      kept=$((kept + 1))
      continue
    fi
    src=""
    if ((${#sources[@]})); then
      src=$(find_in_sources "$(basename "$name")" "$node") || true
    fi
    if [[ -z $src ]]; then
      warn "not found in the driver store: $name"
      continue
    fi
    install -Dm644 "$src" "$stage_dir/$name"
    set_manifest_line "$name" "$(sha256sum "$src" | cut -d' ' -f1)" "$src"
    copied=$((copied + 1))
  done
  write_manifest "$stage_dir/manifest"
  log "staged $copied of ${#missing_names[@]} firmware files in $stage_dir ($kept already staged)"
  exit 0
fi

# ---------------------------------------------------------------------------
# 4. Install: from the stage, else from a Windows partition still on disk
# ---------------------------------------------------------------------------

[[ -n $stage_dir ]] || stage_dir=$DEFAULT_STAGE_DIR
if [[ -n $driverstore ]]; then
  sources=("$driverstore")
else
  # Windows supplies what the stage lacks, so a stage never hides it.
  if [[ -f "$stage_dir/manifest" ]]; then
    log "using the firmware staged by the installer in $stage_dir"
    use_stage "$stage_dir"
  fi
  # After installation, only internal disks are read; -d reads any other.
  if stage_incomplete; then
    scan_windows_partitions internal-only
  fi
fi

mkdir -p "$STATE_DIR"
touch "$MANIFEST"
read_manifest "$MANIFEST"

# A board may list a name after an earlier run installed it. That copy still
# loads first, but this tool never deletes it: matching bytes would not prove
# that the file is still the one it installed, since a board's own tool may
# keep the same file there. Stop recording it and say where it is.
for name in "${provided_names[@]}"; do
  [[ -n ${manifest[$name]-} ]] || continue
  set_manifest_line "$name"
  write_manifest "$MANIFEST"
  file="$FW_ROOT/updates/$name"
  if [[ -e $file || -L $file ]]; then
    log "note: an earlier run installed $file, but $name is now listed in $PROVIDED_DIR/${provided[$name]}.list; the file is left in place and no longer recorded"
    log "note: while it is there it loads before any other copy. The board package may remove it, or you can by hand. This tool does not, because matching bytes would not prove that the file is still the one it installed: a board's own tool may keep the same file there."
  fi
done

installed=0
for i in "${!missing_names[@]}"; do
  node=${missing_nodes[$i]}
  name=${missing_names[$i]}
  src=$(locate_firmware "$name" "$node") || true
  if [[ -z $src ]]; then
    warn "not found: $name"
    continue
  fi
  install -Dm644 "$src" "$FW_ROOT/updates/$name"
  # Record each file as soon as it is in place.
  set_manifest_line "$name" "$(sha256sum "$src" | cut -d' ' -f1)" "$src" \
    "$(date -u +%Y-%m-%dT%H:%M:%SZ)"
  write_manifest "$MANIFEST"
  log "installed $name"
  installed=$((installed + 1))
done

# Include existing GPU zap shaders in the initramfs for the early DRM driver.
declare -a initramfs_files=()
while read -r node name; do
  [[ ${node##*/} == zap-shader || ${node##*/} == zap-shader@* ]] || continue
  if file=$(firmware_file "$name"); then
    initramfs_files+=("${file#"$ROOT"}")
  fi
done < <(dt_firmware_names)
initramfs_changed=0
if ((${#initramfs_files[@]})); then
  initramfs_config=$(
    echo "# Generated by qcom-firmware-extract for early DRM firmware loading."
    printf 'FILES+=(%s)\n' "$(printf '%q ' "${initramfs_files[@]}" | sed 's/ $//')"
  )
  if [[ ! -f $INITCPIO_CONF || $(<"$INITCPIO_CONF") != "$initramfs_config" ]]; then
    mkdir -p "$(dirname "$INITCPIO_CONF")"
    printf '%s\n' "$initramfs_config" >"$INITCPIO_CONF"
    initramfs_changed=1
  fi
elif [[ -f $INITCPIO_CONF ]]; then
  rm -f "$INITCPIO_CONF"
  initramfs_changed=1
fi

still_missing=$(( ${#missing_names[@]} - installed ))
if ((installed == 0)); then
  if ((still_missing > 0 && ${#sources[@]} == 0)); then
    warn "no Windows installation found to copy from; run again with -d <driver store> when one is available"
  fi
  log "nothing installed; $still_missing firmware file(s) still missing"
  ((initramfs_changed)) || exit 0
else
  log "installed $installed firmware file(s), $still_missing still missing"
fi

if [[ $rebuild == yes ]]; then
  if command -v limine-update >/dev/null 2>&1; then
    limine-update
  else
    mkinitcpio -P
  fi
  log "reboot to load the added firmware"
fi