#!/usr/bin/env bash
#
# Purpose : Find the printers on the network and add them as permanent CUPS
#           queues, so they show up in every print dialog. `configure printers`
#           runs this; run it yourself for the options below (--help).
#
# Printers announce themselves over mDNS, and CUPS prints to them through
# temporary queues that only exist while in use. GNOME Settings and Chromium
# browsers list those, but the GTK print dialog ("Print using system
# dialog...") often lists only permanent queues, so a printer can show up in
# Settings and the browser but not there. `lpstat -v` then says "No
# destinations added" while `lpstat -e` lists the printers. A permanent queue
# made with `lpadmin -m everywhere` (driverless IPP Everywhere) fixes that.

set -euo pipefail

COREDESKTOP="$(cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")/.." && pwd)"
source "$COREDESKTOP/configure.d/lib/mdns.sh"

TEST_PAGE=/usr/share/cups/data/default-testpage.pdf

usage() {
  cat << 'EOF'
Usage: setup-printers.sh [options]

Finds the printers on your network and adds the ones you choose as permanent
CUPS queues, so they show up in every print dialog.

  --all           Add every printer found, without asking (a printer another
                  computer shares is left out when it can be reached directly)
  --default NAME  Make queue NAME the default printer afterwards
  --dry-run       Show the commands that would change the system; run none
  --no-deps       Don't install packages or turn on services
  --test-page     Print a test page on each queue added
  -h, --help      Show this help

If a printer is still missing from a print dialog:
  - Quit the app completely (browsers especially) and open it again; print
    dialogs read the list of printers when they start.
  - GTK needs its CUPS print backend: pacman -Ql gtk3 gtk4 | grep printbackend
    If only the "file" backend is there: sudo pacman -S libcups gtk3 gtk4
  - Flatpak apps print through the portal, which needs xdg-desktop-portal and
    xdg-desktop-portal-gnome; the app may also need:
    flatpak override --user --socket=cups <app id>
  - $CUPS_SERVER or ~/.cups/client.conf can point apps at another CUPS server
    than this computer's.
  - To see what the dialog does: GTK_DEBUG=printing <app>, while watching
    journalctl -fu cups
  - cups-browsed can add queues for the printers it finds by itself, instead
    of this script; this one adds only the queues you choose.
EOF
}

DRY_RUN=false ALL=false NO_DEPS=false PRINT_TEST=false DEFAULT=""
while [ $# -gt 0 ]; do
  case "$1" in
    --all) ALL=true ;;
    --default)
      [ $# -ge 2 ] || { echo "--default needs a queue name." >&2; exit 2; }
      DEFAULT="$2"
      shift ;;
    --dry-run) DRY_RUN=true ;;
    --no-deps) NO_DEPS=true ;;
    --test-page) PRINT_TEST=true ;;
    -h|--help) usage; exit 0 ;;
    *) echo "Unknown option: $1 (see --help)" >&2; exit 2 ;;
  esac
  shift
done

# Run a command that changes the system, or with --dry-run only show it
run() {
  if $DRY_RUN; then
    printf '  would run:'
    printf ' %q' "$@"
    echo
  else
    "$@"
  fi
}

# --- Step 1: packages and services ------------------------------------------

setup_dependencies() {
  local pkg service line new backup missing=() stopped=()

  for pkg in cups avahi nss-mdns; do
    pacman -Q "$pkg" &> /dev/null || missing+=("$pkg")
  done
  if [ ${#missing[@]} -gt 0 ]; then
    echo "Installing ${missing[*]}..."
    run sudo pacman -S --needed --noconfirm "${missing[@]}"
  fi

  for service in cups.service avahi-daemon.service; do
    if ! systemctl is-enabled -q "$service" 2> /dev/null || ! systemctl is-active -q "$service"; then
      stopped+=("$service")
    fi
  done
  if [ ${#stopped[@]} -gt 0 ]; then
    echo "Turning on ${stopped[*]}..."
    run sudo systemctl enable --now "${stopped[@]}"
  fi

  # Without mdns_minimal, the .local names in printers' addresses don't resolve
  line="$(grep -m1 '^hosts:' /etc/nsswitch.conf || true)"
  if [[ "$line" != *"mdns_minimal [NOTFOUND=return]"* ]]; then
    if [[ "$line" == *" resolve"* ]]; then
      new="${line/ resolve/ mdns_minimal [NOTFOUND=return] resolve}"
    elif [[ "$line" == *" dns"* ]]; then
      new="${line/ dns/ mdns_minimal [NOTFOUND=return] dns}"
    else
      new="${line:-hosts:} mdns_minimal [NOTFOUND=return]"
    fi
    backup="/etc/nsswitch.conf.bak.$(date +%Y%m%d%H%M%S)"
    echo "Letting .local names resolve, in /etc/nsswitch.conf (a copy is kept in $backup):"
    echo "  before: $line"
    echo "  after:  $new"
    run sudo cp -p /etc/nsswitch.conf "$backup"
    if [ -n "$line" ]; then
      run sudo sed -i "s|^hosts:.*|$new|" /etc/nsswitch.conf
    else
      echo "$new" | run sudo tee -a /etc/nsswitch.conf > /dev/null
    fi
  fi
}

$NO_DEPS || setup_dependencies

if ! LC_ALL=C lpstat -r 2> /dev/null | grep -q "scheduler is running"; then
  echo "CUPS isn't running, so printers can't be found or added." >&2
  echo "Start it with: sudo systemctl enable --now cups.service" >&2
  exit 1
fi

# --- Step 2: find printers ----------------------------------------------------

# %20 -> space, %40 -> @ and any other percent escape
url_decode() {
  local s="${1//\\/\\\\}"
  printf '%b' "${s//%/\\x}"
}

# Printers can take a few seconds to answer
echo "Looking for printers..."
found=""
for attempt in 1 2 3; do
  found="$(LC_ALL=C lpinfo --include-schemes dnssd -v 2> /dev/null | grep '^network dnssd://' || true)"
  [ -n "$found" ] && break
  [ "$attempt" -lt 3 ] && sleep 3
done
if [ -z "$found" ]; then
  echo "No printers found on the network. Is the printer on, and on the same network?"
  exit 0
fi

# What Avahi knows about each printer, by name: its host, address, port and
# path, for spotting queues that already reach it and for the ipp:// fallback
declare -A AVAHI_HOST=() AVAHI_ADDR=() AVAHI_PORT=() AVAHI_RP=()
while IFS='|' read -r name _ host address port txt; do
  [ -n "$name" ] || continue
  AVAHI_HOST["$name"]="$host" AVAHI_ADDR["$name"]="$address" AVAHI_PORT["$name"]="$port"
  if [[ "$txt" =~ \"rp=([^\"]*)\" ]]; then
    AVAHI_RP["$name"]="${BASH_REMATCH[1]}"
  fi
done < <(mdns_services _ipp._tcp || true)

# One entry per printer found: P_NAME is its name as shown, P_BASE the name
# without the "@ host" of a shared printer, P_SHARER that host
P_NAME=() P_BASE=() P_URI=() P_SHARER=()
dnssd='^dnssd://([^/]+)\._ipps?\._tcp\.[^/]*(/[^?]*)'
while read -r _ uri; do
  [[ "$uri" =~ $dnssd ]] || continue
  name="$(url_decode "${BASH_REMATCH[1]}")"
  path="${BASH_REMATCH[2]}"
  sharer=""
  # Another computer's CUPS shares it, as "<printer> @ <computer>"
  if [[ "$path" == /cups* && "$name" == *" @ "* ]]; then
    sharer="${name##* @ }"
  fi
  P_NAME+=("$name") P_BASE+=("${name% @ *}") P_URI+=("$uri") P_SHARER+=("$sharer")
done <<< "$found"
[ "${#P_NAME[@]}" -gt 0 ] || { echo "No printers found on the network."; exit 0; }

# A shared printer that can also be reached directly: the direct one is used,
# so printing doesn't depend on the other computer being on
P_DUP=()
for i in "${!P_NAME[@]}"; do
  P_DUP[i]=""
  [ -n "${P_SHARER[i]}" ] || continue
  for j in "${!P_NAME[@]}"; do
    if [ -z "${P_SHARER[j]}" ] && [ "${P_NAME[j]}" = "${P_BASE[i]}" ]; then
      P_DUP[i]="$j"
    fi
  done
done

# --- Step 3: queue names ------------------------------------------------------

# The queues there are now: name -> device URI
declare -A QUEUE_URI=()
while read -r _ _ queue uri; do
  QUEUE_URI["${queue%:}"]="$uri"
done < <(LC_ALL=C lpstat -v 2> /dev/null | grep '^device for ' || true)

lower() {
  tr '[:upper:]' '[:lower:]' <<< "${1%.}"
}

# The queue that already prints to printer $1, if any: one with the same URI,
# or one made some other way (e.g. GNOME Settings) that reaches the same host
# and path, like ipp://NPI3BE90A.local:631/ipp/print
existing_queue() {
  local i="$1" name="${P_NAME[$1]}" queue uri host rp
  host="$(lower "${AVAHI_HOST[$name]:-}")"
  rp="${AVAHI_RP[$name]:-}"
  for queue in "${!QUEUE_URI[@]}"; do
    uri="${QUEUE_URI[$queue]}"
    if [ "$uri" = "${P_URI[i]}" ]; then
      echo "$queue"
      return 0
    fi
    if [ -n "$host" ] && [[ "$uri" =~ ^ipps?://([^/:]+)(:[0-9]+)?/(.*)$ ]]; then
      local uri_host uri_path="${BASH_REMATCH[3]}"
      uri_host="$(lower "${BASH_REMATCH[1]}")"
      if { [ "$uri_host" = "$host" ] || [ "$uri_host" = "${AVAHI_ADDR[$name]:-}" ]; } \
          && [ "$uri_path" = "$rp" ]; then
        echo "$queue"
        return 0
      fi
    fi
  done
  return 1
}

# "HP LaserJet MFP M234dw (3BE90A) @ lab" -> HP_LaserJet_MFP_M234dw_3BE90A_lab
queue_name() {
  local q="${1// @ /_}"
  q="$(sed -E 's/[^A-Za-z0-9_-]/_/g; s/_+/_/g; s/^_+//; s/_+$//' <<< "$q")"
  echo "${q:0:127}"
}

# A name no other queue has, adding _2, _3... if needed
declare -A TAKEN=()
for queue in "${!QUEUE_URI[@]}"; do TAKEN["$queue"]=1; done
unique_queue_name() {
  local base n=2 name
  base="$(queue_name "$1")"
  name="$base"
  while [ -n "${TAKEN[$name]:-}" ]; do
    name="${base:0:120}_$n"
    n=$((n + 1))
  done
  TAKEN["$name"]=1
  echo "$name"
}

P_EXISTING=()
for i in "${!P_NAME[@]}"; do
  P_EXISTING[i]="$(existing_queue "$i" || true)"
done

# --- Choose ---------------------------------------------------------------------

# Chosen unless told otherwise: printers without a queue that aren't a shared
# copy of a direct one
defaults=()
for i in "${!P_NAME[@]}"; do
  if [ -z "${P_EXISTING[i]}" ] && [ -z "${P_DUP[i]}" ]; then
    defaults+=("$((i + 1))")
  fi
done

echo
echo "Printers on your network:"
for i in "${!P_NAME[@]}"; do
  if [ -n "${P_SHARER[i]}" ]; then how="shared by ${P_SHARER[i]}"; else how="direct"; fi
  note=""
  if [ -n "${P_EXISTING[i]}" ]; then
    note="already set up as ${P_EXISTING[i]}"
  elif [ -n "${P_DUP[i]}" ]; then
    note="same printer as $((P_DUP[i] + 1)), reached directly"
  fi
  printf '  %2d) %-46s %-20s %s\n' "$((i + 1))" "${P_NAME[i]}" "$how" "$note"
done
echo

chosen=()
if $ALL; then
  chosen=("${defaults[@]}")
else
  answer=""
  if [ ${#defaults[@]} -gt 0 ]; then
    read -rp "Add which printers? Numbers separated by spaces [${defaults[*]}], or n for none: " answer || answer=n
    answer="${answer:-${defaults[*]}}"
  else
    read -rp "Every printer is already set up. Add one anyway? Number, or Enter for none: " answer || answer=""
  fi
  for n in ${answer//,/ }; do
    [[ "$n" == [nN]* ]] && { chosen=(); break; }
    if [[ "$n" =~ ^[0-9]+$ ]] && (( n >= 1 && n <= ${#P_NAME[@]} )); then
      chosen+=("$n")
    else
      echo "Skipping \"$n\": not a number from the list."
    fi
  done
fi

# --- Step 4: add queues ---------------------------------------------------------

ADDED=() SKIPPED=() FAILED=()

add_queue() {
  local i="$1" name="${P_NAME[$1]}" queue out fallback=""
  if [ -n "${P_EXISTING[i]}" ]; then
    SKIPPED+=("$name (already set up as ${P_EXISTING[i]})")
    return 0
  fi
  queue="$(unique_queue_name "$name")"
  echo "Adding $name as $queue..."
  if $DRY_RUN; then
    run sudo lpadmin -p "$queue" -E -v "${P_URI[i]}" -m everywhere
    ADDED+=("$queue")
    return 0
  fi

  # lpadmin warns that printer drivers are deprecated; only its status counts
  if out="$(sudo lpadmin -p "$queue" -E -v "${P_URI[i]}" -m everywhere 2>&1)"; then
    ADDED+=("$queue")
    return 0
  fi

  # The printer's address and path from Avahi instead of its mDNS name
  if [ -n "${AVAHI_ADDR[$name]:-}" ]; then
    fallback="ipp://${AVAHI_ADDR[$name]}:${AVAHI_PORT[$name]:-631}/${AVAHI_RP[$name]:-ipp/print}"
    echo "  That didn't work; trying $fallback"
    if out="$(sudo lpadmin -p "$queue" -E -v "$fallback" -m everywhere 2>&1)"; then
      ADDED+=("$queue")
      return 0
    fi
  fi
  unset 'TAKEN[$queue]'
  FAILED+=("$name: $(tail -n1 <<< "$out")")
}

for n in "${chosen[@]}"; do
  add_queue "$((n - 1))"
done

# --- Step 5: default, check and summary ---------------------------------------

if [ -z "$DEFAULT" ] && ! $ALL && ! $DRY_RUN && [ ${#ADDED[@]} -gt 0 ]; then
  current="$(LC_ALL=C lpstat -d 2> /dev/null | sed -n 's/^system default destination: //p')"
  read -rp "Make one of the new printers the default? ${ADDED[*]} (Enter to leave it${current:+ as $current}): " DEFAULT || DEFAULT=""
fi
if [ -n "$DEFAULT" ]; then
  if [ -n "${QUEUE_URI[$DEFAULT]:-}" ] || [[ " ${ADDED[*]} " == *" $DEFAULT "* ]]; then
    run sudo lpadmin -d "$DEFAULT"
  else
    FAILED+=("default printer: there's no queue named $DEFAULT")
    DEFAULT=""
  fi
fi

if ! $DRY_RUN; then
  ready=()
  for queue in "${ADDED[@]}"; do
    if LC_ALL=C lpstat -a "$queue" 2> /dev/null | grep -q "accepting requests"; then
      ready+=("$queue")
      if $PRINT_TEST; then
        lp -d "$queue" "$TEST_PAGE" > /dev/null || FAILED+=("$queue: the test page wasn't sent")
      fi
    else
      FAILED+=("$queue: added, but it isn't accepting jobs (see lpstat -t)")
    fi
  done
  ADDED=("${ready[@]}")
fi

echo
if $DRY_RUN; then echo "Would add:"; else echo "Added:"; fi
[ ${#ADDED[@]} -gt 0 ] && printf '  %s\n' "${ADDED[@]}" || echo "  none"
if [ ${#SKIPPED[@]} -gt 0 ]; then
  echo "Skipped:"
  printf '  %s\n' "${SKIPPED[@]}"
fi
if [ ${#FAILED[@]} -gt 0 ]; then
  echo "Failed:"
  printf '  %s\n' "${FAILED[@]}"
  echo "A printer that can't be added without a driver needs its maker's: hplip for HP,"
  echo "brother-* from the AUR for Brother. This script doesn't install drivers."
fi
[ -n "$DEFAULT" ] && echo "Default printer: $DEFAULT"
if [ ${#ADDED[@]} -gt 0 ] && ! $DRY_RUN; then
  $PRINT_TEST && echo "A test page was sent to each new printer."
  echo "Quit and reopen apps that are open (browsers especially) to see the new printers."
fi

# Without its CUPS backend, the GTK dialog can't list printers at all
if pacman -Q gtk3 &> /dev/null && ! pacman -Ql gtk3 | grep -q 'printbackend-cups'; then
  echo "GTK's CUPS print backend is missing; reinstall it: sudo pacman -S libcups gtk3 gtk4"
fi

[ ${#FAILED[@]} -eq 0 ]