#!/bin/bash

# omarchy:summary=Show or configure the system DNS provider
# omarchy:args=[Cloudflare|Google|DHCP|Custom]
# omarchy:examples=omarchy dns | omarchy dns Cloudflare | omarchy dns Custom

set -euo pipefail

NM_DNS_CONF=/etc/NetworkManager/conf.d/20-omarchy-dns.conf

provider_from_arg() {
  case "${1:-}" in
  Cloudflare | cloudflare)
    echo "Cloudflare"
    ;;
  Google | google)
    echo "Google"
    ;;
  DHCP | dhcp)
    echo "DHCP"
    ;;
  Custom | custom)
    echo "Custom"
    ;;
  *)
    return 1
    ;;
  esac
}

# The path etc/sudoers.d/omarchy-dns names. The privileged half always runs from
# there rather than from whichever copy was invoked, so the rule matches even
# where $OMARCHY_PATH points at a checkout.
PACKAGED_PATH=/usr/bin/omarchy-dns

# True when sudo would run this exact command without stopping for a password.
# `sudo -l` on its own reports whether a command is permitted, which the blanket
# %wheel rule answers yes to for everything; the long listing prints the matched
# entry's tags, so !authenticate is the grant in etc/sudoers.d/omarchy-dns and
# nothing else. Listing runs nothing and, under -n, prompts for nothing, so a
# machine whose omarchy-settings predates that file falls through to polkit
# instead of dying on a password prompt it has no terminal to show.
sudo_grants_passwordless() {
  sudo -n -l -l "$PACKAGED_PATH" "$@" 2>/dev/null | grep -q '!authenticate'
}

require_root() {
  if (( EUID == 0 )); then
    return
  elif [[ -t 0 ]] || sudo_grants_passwordless "$@"; then
    # A terminal can carry sudo's own password prompt. Without one, sudo is
    # right only where the grant reaches; polkit can at least put a prompt on
    # screen, and offer to authenticate as someone else.
    exec sudo "$PACKAGED_PATH" "$@"
  else
    exec pkexec "$PACKAGED_PATH" "$@"
  fi
}

networkmanager_global_dns() {
  [[ -f $NM_DNS_CONF ]] || return 0

  awk -F= '
    /^[[:space:]]*#/ { next }
    /^[[:space:]]*\[global-dns-domain-\*\][[:space:]]*$/ { in_default = 1; next }
    /^[[:space:]]*\[/ { in_default = 0 }
    in_default && /^[[:space:]]*servers[[:space:]]*=/ {
      value = $0
      sub(/^[^=]*=/, "", value)
      print value
      exit
    }
  ' "$NM_DNS_CONF"
}

resolved_dns() {
  awk -F= '
    /^[[:space:]]*#/ { next }
    /^[[:space:]]*DNS[[:space:]]*=/ {
      value=$0
      sub(/^[^=]*=/, "", value)
      print value
      exit
    }
  ' /etc/systemd/resolved.conf 2>/dev/null || true
}

current_dns_provider() {
  local dns=""
  local compact=""

  dns=$(networkmanager_global_dns)
  if [[ -z $(printf '%s' "$dns" | tr -d '[:space:],') ]]; then
    dns=$(resolved_dns)
  fi

  compact=$(printf '%s' "$dns" | tr -d '[:space:],')

  if [[ -z $compact ]]; then
    echo "DHCP"
  elif [[ $dns == *"cloudflare-dns.com"* || $dns == *"1.1.1.1"* || $dns == *"2606:4700:4700::1111"* ]]; then
    echo "Cloudflare"
  elif [[ $dns == *"dns.google"* || $dns == *"8.8.8.8"* || $dns == *"2001:4860:4860::8888"* ]]; then
    echo "Google"
  else
    echo "Custom"
  fi
}

normalize_servers() {
  printf '%s\n' "$*" | tr ',\t\n' ' ' | xargs | tr ' ' ','
}

split_dns_servers() {
  local servers="$1"
  local server clean
  ipv4_dns=""
  ipv6_dns=""

  for server in ${servers//,/ }; do
    clean=${server#dns+tls://}
    clean=${clean#dns+udp://}
    clean=${clean%%#*}
    clean=${clean#[}
    clean=${clean%]}

    [[ -n $clean ]] || continue
    if [[ $clean == *:* ]]; then
      ipv6_dns+="${ipv6_dns:+ }$clean"
    else
      ipv4_dns+="${ipv4_dns:+ }$clean"
    fi
  done
}

write_networkmanager_dns() {
  local servers="$1"

  install -d -m 0755 "$(dirname "$NM_DNS_CONF")"
  cat >"$NM_DNS_CONF" <<EOF
# Managed by omarchy-dns. Remove this file or run omarchy dns DHCP to use DHCP DNS again.
[global-dns]

[global-dns-domain-*]
servers=$servers
EOF
}

clear_networkmanager_dns() {
  rm -f "$NM_DNS_CONF"
}

networkmanager_dns_connection() {
  case "$1" in
    802-11-wireless|802-3-ethernet) return 0 ;;
    *) return 1 ;;
  esac
}

set_connection_dns() {
  local uuid type
  local ipv4_dns="${1:-}"
  local ipv6_dns="${2:-}"

  while IFS=: read -r uuid type; do
    [[ -n $uuid ]] || continue
    networkmanager_dns_connection "$type" || continue

    nmcli connection modify "$uuid" \
      ipv4.ignore-auto-dns yes \
      ipv4.dns "$ipv4_dns" \
      ipv6.ignore-auto-dns yes \
      ipv6.dns "$ipv6_dns" \
      >/dev/null
  done < <(nmcli -t -f UUID,TYPE connection show)
}

clear_connection_dns() {
  local uuid type

  while IFS=: read -r uuid type; do
    [[ -n $uuid ]] || continue
    networkmanager_dns_connection "$type" || continue

    nmcli connection modify "$uuid" \
      ipv4.ignore-auto-dns no \
      ipv4.dns "" \
      ipv6.ignore-auto-dns no \
      ipv6.dns "" \
      >/dev/null
  done < <(nmcli -t -f UUID,TYPE connection show)
}

reapply_active_dns_connections() {
  local device type state

  while IFS=: read -r device type state; do
    [[ -n $device && $state == connected ]] || continue
    case "$type" in
      wifi|ethernet)
        nmcli device reapply "$device" >/dev/null 2>&1 || true
        ;;
    esac
  done < <(nmcli -t -f DEVICE,TYPE,STATE device status)
}

reload_dns_stack() {
  if systemctl is-active --quiet NetworkManager.service 2>/dev/null; then
    # Load the updated NetworkManager config first, then reapply the active
    # profiles. A single conf,dns-full reload here pushes the old active DNS
    # settings, making the shell toggle appear one selection behind.
    nmcli general reload conf >/dev/null 2>&1 || systemctl reload NetworkManager.service 2>/dev/null || true
    reapply_active_dns_connections
  fi

  systemctl reload systemd-resolved.service 2>/dev/null || systemctl restart systemd-resolved.service

  if systemctl is-active --quiet NetworkManager.service 2>/dev/null; then
    # A resolved reload/restart can leave per-link DNS stale or empty; ask
    # NetworkManager to publish DNS after resolved has reread its config.
    nmcli general reload dns-full >/dev/null 2>&1 || true
  fi
}

usage() {
  echo "Usage: omarchy-dns [Cloudflare|Google|DHCP|Custom]" >&2
}

if (( $# == 0 )); then
  current_dns_provider
  exit 0
fi

if (( $# > 1 )); then
  usage
  exit 1
fi

if ! provider=$(provider_from_arg "$1"); then
  usage
  exit 1
fi

require_root "$provider"

case "$provider" in
Cloudflare)
  write_networkmanager_dns "1.1.1.1,1.0.0.1,2606:4700:4700::1111,2606:4700:4700::1001"
  set_connection_dns "1.1.1.1 1.0.0.1" "2606:4700:4700::1111 2606:4700:4700::1001"
  tee /etc/systemd/resolved.conf >/dev/null <<'EOF'
[Resolve]
DNS=1.1.1.1#cloudflare-dns.com 1.0.0.1#cloudflare-dns.com 2606:4700:4700::1111#cloudflare-dns.com 2606:4700:4700::1001#cloudflare-dns.com
FallbackDNS=9.9.9.9#dns.quad9.net 149.112.112.112#dns.quad9.net 2620:fe::fe#dns.quad9.net 2620:fe::9#dns.quad9.net
DNSOverTLS=opportunistic
EOF
  ;;

Google)
  write_networkmanager_dns "8.8.8.8,8.8.4.4,2001:4860:4860::8888,2001:4860:4860::8844"
  set_connection_dns "8.8.8.8 8.8.4.4" "2001:4860:4860::8888 2001:4860:4860::8844"
  tee /etc/systemd/resolved.conf >/dev/null <<'EOF'
[Resolve]
DNS=8.8.8.8#dns.google 8.8.4.4#dns.google 2001:4860:4860::8888#dns.google 2001:4860:4860::8844#dns.google
FallbackDNS=9.9.9.9#dns.quad9.net 149.112.112.112#dns.quad9.net 2620:fe::fe#dns.quad9.net 2620:fe::9#dns.quad9.net
DNSOverTLS=opportunistic
EOF
  ;;

DHCP)
  clear_networkmanager_dns
  clear_connection_dns
  tee /etc/systemd/resolved.conf >/dev/null <<'EOF'
[Resolve]
DNSOverTLS=no
EOF
  ;;

Custom)
  echo "Enter your DNS servers (space-separated, e.g. '192.168.1.1 1.1.1.1'):"
  if ! read -r dns_servers; then
    dns_servers=""
  fi

  dns_servers=$(normalize_servers "$dns_servers")
  if [[ -z $dns_servers ]]; then
    echo "Error: No DNS servers provided." >&2
    exit 1
  fi

  split_dns_servers "$dns_servers"
  write_networkmanager_dns "$dns_servers"
  set_connection_dns "$ipv4_dns" "$ipv6_dns"
  tee /etc/systemd/resolved.conf >/dev/null <<EOF
[Resolve]
DNS=${dns_servers//,/ }
FallbackDNS=9.9.9.9#dns.quad9.net 149.112.112.112#dns.quad9.net 2620:fe::fe#dns.quad9.net 2620:fe::9#dns.quad9.net
EOF
  ;;
esac

reload_dns_stack