#!/bin/sh
set -eu

RAM_ROOT="${HBP_RAM_ROOT:-/run/hbp-ram-runtime}"
MIN_AVAIL_MB="${HBP_RAM_MIN_AVAIL_MB:-100}"
SOCK="${CUPS_SERVER_SOCK:-/var/run/cups/cups.sock}"
TMP_RAW="/tmp/hbp-ram-roots.raw.$$"
TMP_LIST="/tmp/hbp-ram-roots.list.$$"
TMP_DONE_FILES="/tmp/hbp-ram-done-files.$$"
trap 'rm -f "$TMP_RAW" "$TMP_LIST" "$TMP_DONE_FILES"' EXIT INT TERM

usage() {
  cat <<USAGE
usage:
  cups-runtime-ram-feasibility estimate [core|full]
  cups-runtime-ram-feasibility stage [core|full]
  cups-runtime-ram-feasibility check
  cups-runtime-ram-feasibility status
  cups-runtime-ram-feasibility detach-sd
  cups-runtime-ram-feasibility unstage

notes:
  - core: curated runtime roots for HBP print path.
  - full: all paths from /cups-runtime-store-paths (requires latest image build).
  - env:
      HBP_RAM_ROOT=/run/hbp-ram-runtime
      HBP_RAM_SIZE=<tmpfs size, e.g. 320m>   (stage only)
      HBP_RAM_MIN_AVAIL_MB=100               (check gate)
USAGE
}

store_root_of() {
  p="${1:-}"
  [ -n "$p" ] || return 1
  rp="$(readlink -f "$p" 2>/dev/null || true)"
  [ -n "$rp" ] || return 1
  printf '%s\n' "$rp" | sed -n 's#^\(/nix/store/[^/]*-[^/]*\).*#\1#p' | head -n 1
}

add_root() {
  r="${1:-}"
  [ -n "$r" ] || return 0
  case "$r" in
    /nix/store/*)
      [ -e "$r" ] && echo "$r" >> "$TMP_RAW"
      case "$r" in
        *-lib) ;;
        *)
          # Nix frequently splits runtime libraries into sibling -lib outputs.
          if [ -e "${r}-lib" ]; then
            echo "${r}-lib" >> "$TMP_RAW"
          fi
          ;;
      esac
      ;;
  esac
}

mem_field_mb() {
  key="$1"
  awk -v k="$key" '$1==k":" {printf "%d", $2/1024; exit}' /proc/meminfo
}

report_mem() {
  tag="$1"
  echo "mem[$tag]: total=$(mem_field_mb MemTotal)MB avail=$(mem_field_mb MemAvailable)MB free=$(mem_field_mb MemFree)MB"
}

resolve_lib_path() {
  lib="$1"
  runpaths="$2"
  oldifs="$IFS"
  IFS=':'
  for d in $runpaths /lib /usr/lib; do
    [ -n "$d" ] || continue
    if [ -e "$d/$lib" ]; then
      IFS="$oldifs"
      readlink -f "$d/$lib" 2>/dev/null || true
      return 0
    fi
  done
  IFS="$oldifs"

  # BusyBox find behavior can vary; keep fallback lookup simple and robust.
  p="$(find /nix/store -name "$lib" 2>/dev/null | head -n 1 || true)"
  if [ -n "$p" ]; then
    readlink -f "$p" 2>/dev/null || true
    return 0
  fi
  return 1
}

collect_loader_roots() {
  f="$1"
  [ -n "$f" ] || return 0
  [ -x /lib/ld-musl-armhf.so.1 ] || return 0

  /lib/ld-musl-armhf.so.1 --list "$f" 2>/dev/null \
    | awk '{for(i=1;i<=NF;i++) if($i ~ /^\/nix\/store\//) print $i}' \
    | while read -r p; do
        add_root "$(store_root_of "$p" || true)"
      done
}

collect_file_deps() {
  f="$1"
  [ -n "$f" ] || return 0
  f="$(readlink -f "$f" 2>/dev/null || true)"
  [ -n "$f" ] || return 0
  [ -e "$f" ] || return 0

  grep -Fxq "$f" "$TMP_DONE_FILES" 2>/dev/null && return 0
  echo "$f" >> "$TMP_DONE_FILES"
  add_root "$(store_root_of "$f" || true)"
  collect_loader_roots "$f"

  [ -x /bin/readelf ] || return 0

  interp="$(/bin/readelf -l "$f" 2>/dev/null | sed -n 's@.*Requesting program interpreter: \(.*\)]@\1@p' | head -n 1 || true)"
  [ -n "$interp" ] && collect_file_deps "$interp"

  runpaths="$(/bin/readelf -d "$f" 2>/dev/null | sed -n 's@.*Library runpath: \[\(.*\)\]@\1@p' | head -n 1 || true)"
  if [ -n "$runpaths" ]; then
    oldifs="$IFS"
    IFS=':'
    for rp in $runpaths; do
      case "$rp" in
        /nix/store/*) add_root "$(store_root_of "$rp" || true)" ;;
      esac
    done
    IFS="$oldifs"
  fi

  needed="$(/bin/readelf -d "$f" 2>/dev/null | sed -n 's@.*Shared library: \[\(.*\)\]@\1@p' || true)"
  [ -n "$needed" ] || return 0
  for lib in $needed; do
    libp="$(resolve_lib_path "$lib" "$runpaths" || true)"
    [ -n "$libp" ] || continue
    collect_file_deps "$libp"
  done
}

collect_core_roots() {
  : > "$TMP_RAW"
  : > "$TMP_DONE_FILES"
  for cmd in cupsd lp lpadmin lpstat lpinfo ppdc gs pdftops cupsfilter; do
    p="$(command -v "$cmd" 2>/dev/null || true)"
    [ -n "$p" ] || continue
    rp="$(readlink -f "$p" 2>/dev/null || true)"
    [ -n "$rp" ] && collect_file_deps "$rp"
  done
  for p in \
    /var/cups-serverbin/lib/cups/filter/rastertobrlaser \
    /var/cups-serverbin/lib/cups/backend/usb \
    /var/cups-serverbin/lib/cups/driver/cups-driverd \
    /var/cups-serverbin/lib/cups/driver/drv
  do
    [ -e "$p" ] || continue
    rp="$(readlink -f "$p" 2>/dev/null || true)"
    [ -n "$rp" ] && collect_file_deps "$rp"
  done
  if [ -f /cups-runtime.env ]; then
    # shellcheck source=/dev/null
    . /cups-runtime.env
    add_root "${BRLASER_ROOT:-}"
    add_root "${CUPS_FILTERS_ROOT:-}"
  fi
  cupsd_real="$(readlink -f /bin/cupsd 2>/dev/null || true)"
  cupsd_root="$(store_root_of "$cupsd_real" || true)"
  add_root "$cupsd_root"
  if [ -n "$cupsd_root" ] && [ -d "${cupsd_root}-lib" ]; then
    add_root "${cupsd_root}-lib"
  fi
  if [ -f /etc/cups/cups-files.conf ]; then
    awk '{for(i=1;i<=NF;i++) if($i ~ /^\/nix\/store\//) print $i}' /etc/cups/cups-files.conf \
      | while read -r p; do
          add_root "$(store_root_of "$p" || true)"
        done
  fi
  sort -u "$TMP_RAW" > "$TMP_LIST"
}

collect_full_roots() {
  : > "$TMP_RAW"
  if [ ! -f /cups-runtime-store-paths ]; then
    echo "error: /cups-runtime-store-paths missing (rebuild with latest flake/initramfs)" >&2
    exit 1
  fi
  while read -r p; do
    [ -n "$p" ] || continue
    case "$p" in
      /nix/store/*) add_root "$p" ;;
    esac
  done < /cups-runtime-store-paths
  sort -u "$TMP_RAW" > "$TMP_LIST"
}

estimate_kib() {
  total=0
  while read -r r; do
    [ -e "$r" ] || continue
    sz="$(du -sk "$r" 2>/dev/null | awk '{print $1}')"
    [ -n "$sz" ] || sz=0
    total=$((total + sz))
  done < "$TMP_LIST"
  echo "$total"
}

default_tmpfs_size() {
  kib="$1"
  # +20% copy overhead and +64 MiB headroom.
  size_kib=$(( (kib * 120) / 100 + 65536 ))
  echo "${size_kib}k"
}

is_mounted() {
  mnt="$1"
  awk -v m="$mnt" '$2==m{found=1} END{exit found?0:1}' /proc/mounts
}

setup_tmpfs() {
  size_opt="$1"
  mkdir -p "$RAM_ROOT"
  if ! is_mounted "$RAM_ROOT"; then
    mount -t tmpfs -o "size=$size_opt,mode=0755" tmpfs "$RAM_ROOT"
  fi
  mkdir -p "$RAM_ROOT/nix/store"
}

copy_roots() {
  n=0
  while read -r r; do
    [ -e "$r" ] || continue
    n=$((n + 1))
    dest="$RAM_ROOT/nix/store/$(basename "$r")"
    if [ -e "$dest" ]; then
      src_real="$(readlink -f "$r" 2>/dev/null || true)"
      dst_real="$(readlink -f "$dest" 2>/dev/null || true)"
      if [ -n "$src_real" ] && [ "$src_real" = "$dst_real" ]; then
        echo "copy[$n]: $r (already staged)"
        continue
      fi
    fi
    echo "copy[$n]: $r"
    cp -a "$r" "$RAM_ROOT/nix/store/"
  done < "$TMP_LIST"
}

bind_nix() {
  src="$(awk '$2=="/nix"{print $1}' /proc/mounts | tail -n 1)"
  if [ "$src" = "$RAM_ROOT/nix" ]; then
    return 0
  fi
  mount --bind "$RAM_ROOT/nix" /nix
}

unbind_nix() {
  src="$(awk '$2=="/nix"{print $1}' /proc/mounts | tail -n 1)"
  if [ "$src" = "$RAM_ROOT/nix" ]; then
    umount /nix
  fi
}

teardown_tmpfs() {
  if is_mounted "$RAM_ROOT"; then
    umount "$RAM_ROOT"
  fi
}

status() {
  src="$(awk '$2=="/nix"{print $1 " (" $3 ")"}' /proc/mounts | tail -n 1)"
  [ -n "$src" ] || src="(unmounted)"
  echo "nix-mount: $src"
  if is_mounted "$RAM_ROOT"; then
    rsrc="$(awk -v m="$RAM_ROOT" '$2==m{print $1 " (" $3 ")"}' /proc/mounts | tail -n 1)"
    echo "ram-root: $rsrc"
  else
    echo "ram-root: not-mounted"
  fi
  report_mem "status"
}

enforce_mem_gate() {
  avail="$(mem_field_mb MemAvailable)"
  if [ "$avail" -lt "$MIN_AVAIL_MB" ]; then
    echo "FAIL: MemAvailable=${avail}MB < ${MIN_AVAIL_MB}MB gate" >&2
    return 1
  fi
  echo "PASS: MemAvailable=${avail}MB >= ${MIN_AVAIL_MB}MB gate"
}

summarize_estimate() {
  mode="$1"
  count="$(wc -l < "$TMP_LIST" | tr -d ' ')"
  kib="$(estimate_kib)"
  mib=$((kib / 1024))
  echo "mode=$mode roots=$count size=${mib}MiB"
}

run_check() {
  status
  enforce_mem_gate
  echo "[queues]"
  lpstat -h "$SOCK" -p -v || true
}

detach_sd() {
  src="$(awk '$2=="/nix"{print $1}' /proc/mounts | tail -n 1)"
  fstype="$(awk '$2=="/nix"{print $3}' /proc/mounts | tail -n 1)"
  if [ "$src" != "$RAM_ROOT/nix" ] && [ "$fstype" != "tmpfs" ]; then
    echo "error: /nix is not RAM-backed; run 'stage' first (src=$src fstype=$fstype)" >&2
    return 1
  fi

  sync
  if command -v blockdev >/dev/null 2>&1 && [ -b /dev/mmcblk0 ]; then
    blockdev --flushbufs /dev/mmcblk0 >/dev/null 2>&1 || true
  fi

  pass=0
  while [ "$pass" -lt 6 ]; do
    pass=$((pass + 1))
    remain="$(awk '$1 ~ /^\/dev\/mmcblk0p[0-9]+$/ {print $1 " " $2}' /proc/mounts)"
    [ -z "$remain" ] && break

    # Handle stacked /nix mounts: if top /nix is RAM-backed but an mmc mount
    # still exists at /nix, drop top layer once to expose lower mount.
    if echo "$remain" | awk '$2=="/nix"{found=1} END{exit found?0:1}'; then
      src="$(awk '$2=="/nix"{print $1}' /proc/mounts | tail -n 1)"
      fstype="$(awk '$2=="/nix"{print $3}' /proc/mounts | tail -n 1)"
      if [ "$src" = "$RAM_ROOT/nix" ] || [ "$fstype" = "tmpfs" ]; then
        umount /nix >/dev/null 2>&1 || true
      fi
    fi

    remain="$(awk '$1 ~ /^\/dev\/mmcblk0p[0-9]+$/ {print $1 " " $2}' /proc/mounts)"
    [ -z "$remain" ] && break
    tmp_mmc="/tmp/hbp-ram-mmc.$$.${pass}"
    echo "$remain" > "$tmp_mmc"
    while read -r dev mnt; do
      [ -n "$dev" ] || continue
      [ -n "$mnt" ] || continue
      echo "detach: unmounting $dev ($mnt)"
      umount "$mnt" >/dev/null 2>&1 || umount -l "$mnt" >/dev/null 2>&1 || \
        echo "warn: failed to unmount $dev ($mnt)" >&2
    done < "$tmp_mmc"
    rm -f "$tmp_mmc"

    # Ensure /nix is RAM-backed after detach steps.
    src="$(awk '$2=="/nix"{print $1}' /proc/mounts | tail -n 1)"
    fstype="$(awk '$2=="/nix"{print $3}' /proc/mounts | tail -n 1)"
    if [ "$src" != "$RAM_ROOT/nix" ] && [ "$fstype" != "tmpfs" ] && [ -d "$RAM_ROOT/nix/store" ]; then
      mount --bind "$RAM_ROOT/nix" /nix >/dev/null 2>&1 || true
    fi
  done

  sync
  remain="$(awk '$1 ~ /^\/dev\/mmcblk0p[0-9]+$/ {print $1 " -> " $2}' /proc/mounts)"
  if [ -n "$remain" ]; then
    echo "FAIL: mmc partitions still mounted:" >&2
    echo "$remain" >&2
    return 1
  fi
  echo "SD detach prep complete: no mmcblk0p* mounts remain."
}

nix_is_ram_backed() {
  src="$(awk '$2=="/nix"{print $1}' /proc/mounts | tail -n 1)"
  fstype="$(awk '$2=="/nix"{print $3}' /proc/mounts | tail -n 1)"
  if [ "$src" = "$RAM_ROOT/nix" ] || [ "$fstype" = "tmpfs" ]; then
    return 0
  fi
  if [ -d /nix/store ] && [ -d "$RAM_ROOT/nix/store" ]; then
    nix_store="$(readlink -f /nix/store 2>/dev/null || true)"
    ram_store="$(readlink -f "$RAM_ROOT/nix/store" 2>/dev/null || true)"
    if [ -n "$nix_store" ] && [ "$nix_store" = "$ram_store" ]; then
      return 0
    fi
  fi
  return 1
}

ACTION="${1:-}"
MODE="${2:-core}"

case "$ACTION" in
  estimate)
    case "$MODE" in
      core) collect_core_roots ;;
      full) collect_full_roots ;;
      *) usage; exit 2 ;;
    esac
    summarize_estimate "$MODE"
    ;;
  stage)
    case "$MODE" in
      core) collect_core_roots ;;
      full) collect_full_roots ;;
      *) usage; exit 2 ;;
    esac
    summarize_estimate "$MODE"
    if nix_is_ram_backed; then
      echo "stage: /nix already RAM-backed; skipping copy"
      report_mem "after-stage"
      run_check
      exit 0
    fi
    kib="$(estimate_kib)"
    size_opt="${HBP_RAM_SIZE:-$(default_tmpfs_size "$kib")}"
    report_mem "before-stage"
    setup_tmpfs "$size_opt"
    copy_roots
    bind_nix
    report_mem "after-stage"
    run_check
    ;;
  check)
    run_check
    ;;
  status)
    status
    ;;
  detach-sd)
    detach_sd
    ;;
  unstage)
    unbind_nix
    teardown_tmpfs
    report_mem "after-unstage"
    ;;
  *)
    usage
    exit 2
    ;;
esac
