inventory.nix

@@ -64,7 +64,7 @@
   #   "external" = {
   #     devices."remote" = builtins.fetchGit {            # External module via Git
   #       url = "https://github.com/example/config";
-  #       rev = "abc123...";
+  #       rev = "e1ccd7cc3e709afe4f50b0627e1c4bde49165014";
   #     };
   #   };  # ========== Lab Laptops ==========
   # Creates: nix-laptop1, nix-laptop2
@@ -121,8 +121,9 @@
         };
       };
       "usda-dash" = builtins.fetchGit {
-        url = "https://git.factory.uga.edu/MODEL/usda-dash-config.git";
-        rev = "6c0029057aa50d0b4d6a0205c3ded890eb08979c";
+        url = "git@factory.uga.edu:MODEL/usda-dash-config.git";
+        rev = "49cded91cff4a956d4e01ac6b8fe4efa86f82182";
+        submodules = true;
       };
     };
     overrides = {
@@ -150,7 +151,7 @@
   #     # Option 1: fetchGit with specific revision (recommended for reproducibility)
   #     "prod-server" = builtins.fetchGit {
   #       url = "https://github.com/example/server-config";
-  #       rev = "abc123def456...";  # Full commit hash
+  #       rev = "e1ccd7cc3e709afe4f50b0627e1c4bde49165014";  # Full commit hash
   #       ref = "main";              # Optional: branch/tag name
   #     };
   #

sw/builders/services.nix

@@ -18,7 +18,7 @@ mkIf builderCfg.giteaRunner.enable {
     tokenFile = builderCfg.giteaRunner.tokenFile;
     labels = builderCfg.giteaRunner.extraLabels;
     name = builderCfg.giteaRunner.name;
-    
+
     # Run as engr-ugaif user to access SSH keys
     settings = {
       runner = {
@@ -38,7 +38,7 @@ mkIf builderCfg.giteaRunner.enable {
       # Run as engr-ugaif user
       User = mkForce "engr-ugaif";
       Group = mkForce "users";
-      
+
       # Give the service more time to stop cleanly
       TimeoutStopSec = mkForce 60;
 

sw/default.nix

@@ -29,6 +29,7 @@ in
     ./python.nix
     ./ghostty.nix
     ./updater.nix
+    ./update-ref.nix
   ];
 
   options.athenix.sw = {

sw/update-ref.nix

@@ -0,0 +1,487 @@
+{ pkgs, ... }:
+{
+  environment.systemPackages = with pkgs; [
+    python3
+    git
+    (pkgs.writeShellScriptBin "update-ref" ''
+      set -euo pipefail
+
+      RED='\033[31m'; YEL='\033[33m'; NC='\033[0m'
+      die() { printf "''${RED}error:''${NC} %s\n" "$*" >&2; exit 2; }
+      warn() { printf "''${YEL}warning:''${NC} %s\n" "$*" >&2; }
+
+      usage() {
+        cat >&2 <<'EOF'
+      usage:
+        update-ref [-R PATH|--athenix-repo=PATH] [-b BRANCH|--athenix-branch=BRANCH]
+                  [-m "msg"|--message "msg"]
+                  [-p[=false] [remote[=URL]]|--push[=false] [remote[=URL]]]
+                  [--make-local|-l] [--make-remote|-r]
+                  user=<username> | system=<device-type>:<hostkey>
+      EOF
+        exit 2
+      }
+
+      # --- must be in a git repo (current dir) ---
+      git rev-parse --is-inside-work-tree >/dev/null 2>&1 || die "This directory is not a git project"
+      CUR_REPO_ROOT="$(git rev-parse --show-toplevel)"
+      CUR_BRANCH="$(git rev-parse --abbrev-ref HEAD)"
+
+      # --- athenix checkout (working tree) ---
+      ATHENIX_DIR="$HOME/athenix"
+      ATHENIX_BRANCH=""
+
+      # --- current repo automation ---
+      COMMIT_MSG=""
+      PUSH_SPEC=""
+
+      # --- push / url mode ---
+      PUSH_SET=0
+      DO_PUSH=0
+      MODE_FORCE=""  # "", local, remote
+
+      TARGET=""
+
+      is_remote_url() {
+        # https://, http://, ssh://, or scp-style git@host:org/repo
+        printf "%s" "$1" | grep -qE '^(https?|ssh)://|^[^/@:]+@[^/:]+:'
+      }
+
+      derive_full_hostname() {
+        devtype="$1"; hostkey="$2"
+        if printf "%s" "$hostkey" | grep -q '-' || printf "%s" "$hostkey" | grep -q "^$devtype"; then
+          printf "%s" "$hostkey"
+        elif printf "%s" "$hostkey" | grep -qE '^[0-9]+$'; then
+          printf "%s" "$devtype$hostkey"
+        else
+          printf "%s" "$devtype-$hostkey"
+        fi
+      }
+
+      extract_existing_fetch_url() {
+        # args: mode file username key
+        python3 - "$1" "$2" "$3" "$4" <<'PY'
+      import sys, re, pathlib
+      mode, file, username, key = sys.argv[1:5]
+      t = pathlib.Path(file).read_text()
+
+      def url_from_block(block: str) -> str:
+          if not block:
+              return ""
+          m = re.search(r'url\s*=\s*"([^"]+)"\s*;', block)
+          return m.group(1) if m else ""
+
+      if mode == "user":
+          m = re.search(r'(?s)\n\s*' + re.escape(username) + r'\.external\s*=\s*builtins\.fetchGit\s*\{(.*?)\n\s*\};', t)
+          block = m.group(1) if m else ""
+          print(url_from_block(block))
+      else:
+          m = re.search(r'(?s)\n\s*"' + re.escape(key) + r'"\s*=\s*builtins\.fetchGit\s*\{(.*?)\n\s*\};', t)
+          block = m.group(1) if m else ""
+          print(url_from_block(block))
+      PY
+      }
+
+      # --- parse args ---
+      while [ "$#" -gt 0 ]; do
+        case "$1" in
+          user=*|system=*)
+            [ -z "$TARGET" ] || die "Only one subcommand allowed (user=... or system=...)"
+            TARGET="$1"; shift
+            ;;
+          --athenix-repo=*)
+            ATHENIX_DIR="''${1#*=}"; shift
+            ;;
+          -R)
+            [ "$#" -ge 2 ] || usage
+            ATHENIX_DIR="$2"; shift 2
+            ;;
+          --athenix-branch=*)
+            ATHENIX_BRANCH="''${1#*=}"; shift
+            ;;
+          -b)
+            [ "$#" -ge 2 ] || usage
+            ATHENIX_BRANCH="$2"; shift 2
+            ;;
+          -m|--message)
+            [ "$#" -ge 2 ] || usage
+            COMMIT_MSG="$2"; shift 2
+            ;;
+
+          -p|--push)
+            PUSH_SET=1
+            DO_PUSH=1
+            PUSH_SPEC=""
+
+            # If there is a next token, only consume it if it is a remote spec
+            # and not another flag or the subcommand.
+            if [ "$#" -ge 2 ]; then
+              nxt="$2"
+
+              if printf "%s" "$nxt" | grep -qE '^(user=|system=)'; then
+                # next token is the subcommand; don't consume it
+                shift
+              elif printf "%s" "$nxt" | grep -qE '^-'; then
+                # next token is another flag; don't consume it
+                shift
+              elif printf "%s" "$nxt" | grep -qE '^[A-Za-z0-9._-]+$'; then
+                # remote name
+                PUSH_SPEC="$nxt"
+                shift 2
+              elif printf "%s" "$nxt" | grep -qE '^[A-Za-z0-9._-]+=.+$'; then
+                # remote=URL
+                PUSH_SPEC="$nxt"
+                shift 2
+              else
+                # unknown token; treat as not-a-push-spec and don't consume it
+                shift
+              fi
+            else
+              shift
+            fi
+            ;;
+
+          -p=*|--push=*)
+            PUSH_SET=1
+            val="''${1#*=}"
+            case "$val" in
+              false|0|no|off) DO_PUSH=0 ;;
+              true|1|yes|on|"") DO_PUSH=1 ;;
+              *) die "Invalid value for --push: $val (use true/false)" ;;
+            esac
+            shift
+            ;;
+
+          --make-local|-l) MODE_FORCE="local"; shift ;;
+          --make-remote|-r) MODE_FORCE="remote"; shift ;;
+          -h|--help) usage ;;
+          *) die "Unknown argument: $1" ;;
+        esac
+      done
+
+      [ -n "$TARGET" ] || die "Missing required subcommand: user=<username> or system=<device-type>:<hostkey>"
+
+      # --- validate athenix working tree path ---
+      [ -d "$ATHENIX_DIR" ] || die "$ATHENIX_DIR does not exist"
+      git -C "$ATHENIX_DIR" rev-parse --is-inside-work-tree >/dev/null 2>&1 || die "$ATHENIX_DIR is not a git project (athenix checkout)"
+
+      # --- -b behavior: fork/switch athenix working tree into branch ---
+      if [ -n "$ATHENIX_BRANCH" ]; then
+        ATH_CUR_BRANCH="$(git -C "$ATHENIX_DIR" rev-parse --abbrev-ref HEAD)"
+        if [ "$ATH_CUR_BRANCH" != "$ATHENIX_BRANCH" ]; then
+          if git -C "$ATHENIX_DIR" show-ref --verify --quiet "refs/heads/$ATHENIX_BRANCH"; then
+            warn "Branch '$ATHENIX_BRANCH' already exists in $ATHENIX_DIR."
+            warn "Delete and recreate it from current branch '$ATH_CUR_BRANCH' state? [y/N] "
+            read -r ans || true
+            case "''${ans:-N}" in
+              y|Y|yes|YES)
+                git -C "$ATHENIX_DIR" branch -D "$ATHENIX_BRANCH"
+                git -C "$ATHENIX_DIR" switch -c "$ATHENIX_BRANCH"
+                ;;
+              *)
+                git -C "$ATHENIX_DIR" switch "$ATHENIX_BRANCH"
+                ;;
+            esac
+          else
+            git -C "$ATHENIX_DIR" switch -c "$ATHENIX_BRANCH"
+          fi
+        fi
+      fi
+
+      # --- target file + identifiers ---
+      MODE=""; FILE=""; USERNAME=""; DEVTYPE=""; HOSTKEY=""
+      case "$TARGET" in
+        user=*)
+          MODE="user"
+          USERNAME="''${TARGET#user=}"
+          [ -n "$USERNAME" ] || die "user=<username>: username missing"
+          FILE="$ATHENIX_DIR/users.nix"
+          ;;
+        system=*)
+          MODE="system"
+          RHS="''${TARGET#system=}"
+          printf "%s" "$RHS" | grep -q ':' || die "system=... must be system=<device-type>:<hostkey>"
+          DEVTYPE="''${RHS%%:*}"
+          HOSTKEY="''${RHS#*:}"
+          [ -n "$DEVTYPE" ] || die "system=<device-type>:<hostkey>: device-type missing"
+          [ -n "$HOSTKEY" ] || die "system=<device-type>:<hostkey>: hostkey missing"
+          FILE="$ATHENIX_DIR/inventory.nix"
+          ;;
+      esac
+      [ -f "$FILE" ] || die "File not found: $FILE"
+
+      # --- push default based on existing entry url in the target file ---
+      EXISTING_URL=""
+      ENTRY_EXISTS=0
+      if [ "$MODE" = "user" ]; then
+        EXISTING_URL="$(extract_existing_fetch_url user "$FILE" "$USERNAME" "")"
+        [ -n "$EXISTING_URL" ] && ENTRY_EXISTS=1 || true
+      else
+        FULL="$(derive_full_hostname "$DEVTYPE" "$HOSTKEY")"
+        EXISTING_URL="$(extract_existing_fetch_url system "$FILE" "" "$HOSTKEY")"
+        if [ -n "$EXISTING_URL" ]; then
+          ENTRY_EXISTS=1
+        elif [ "$FULL" != "$HOSTKEY" ]; then
+          EXISTING_URL="$(extract_existing_fetch_url system "$FILE" "" "$FULL")"
+          [ -n "$EXISTING_URL" ] && ENTRY_EXISTS=1 || true
+        fi
+      fi
+
+      if [ "$PUSH_SET" -eq 0 ]; then
+        if [ "$ENTRY_EXISTS" -eq 1 ] && is_remote_url "$EXISTING_URL"; then
+          DO_PUSH=1
+        else
+          DO_PUSH=0
+          [ "$MODE_FORCE" = "remote" ] && DO_PUSH=1 || true
+        fi
+      fi
+      if [ "$MODE_FORCE" = "local" ] && [ "$PUSH_SET" -eq 0 ]; then
+        DO_PUSH=0
+      fi
+
+      # --- if current repo dirty, prompt ---
+      if [ -n "$(git status --porcelain)" ]; then
+        warn "This branch has untracked or uncommitted changes. Would you like to add, commit''${DO_PUSH:+, and push}? [y/N] "
+        read -r ans || true
+        case "''${ans:-N}" in
+          y|Y|yes|YES)
+            git add -A
+            if ! git diff --cached --quiet; then
+              if [ -n "$COMMIT_MSG" ]; then git commit -m "$COMMIT_MSG"; else git commit; fi
+            else
+              warn "No staged changes to commit."
+            fi
+            ;;
+          *) warn "Proceeding without committing. (rev will be last committed HEAD.)" ;;
+        esac
+      fi
+
+      # --- push current repo if requested ---
+      PUSH_REMOTE_URL=""
+      if [ "$DO_PUSH" -eq 1 ]; then
+        if [ -n "$PUSH_SPEC" ]; then
+          if printf "%s" "$PUSH_SPEC" | grep -q '='; then
+            REM_NAME="''${PUSH_SPEC%%=*}"
+            REM_URL="''${PUSH_SPEC#*=}"
+            [ -n "$REM_NAME" ] || die "--push remote-name=URL: remote-name missing"
+            [ -n "$REM_URL" ] || die "--push remote-name=URL: URL missing"
+            if git remote get-url "$REM_NAME" >/dev/null 2>&1; then
+              git remote set-url "$REM_NAME" "$REM_URL"
+            else
+              git remote add "$REM_NAME" "$REM_URL"
+            fi
+            git push -u "$REM_NAME" "$CUR_BRANCH"
+            PUSH_REMOTE_URL="$REM_URL"
+          else
+            REM_NAME="$PUSH_SPEC"
+            git push -u "$REM_NAME" "$CUR_BRANCH"
+            PUSH_REMOTE_URL="$(git remote get-url "$REM_NAME")"
+          fi
+        else
+          if ! git rev-parse --abbrev-ref --symbolic-full-name @{u} >/dev/null 2>&1; then
+            die "No upstream is set. Set a default upstream with \"git branch -u <remote>/<remote_branch_name>\""
+          fi
+          git push
+          UPSTREAM_REMOTE="$(git rev-parse --abbrev-ref --symbolic-full-name @{u} | cut -d/ -f1)"
+          PUSH_REMOTE_URL="$(git remote get-url "$UPSTREAM_REMOTE")"
+        fi
+      fi
+
+      CUR_REV="$(git -C "$CUR_REPO_ROOT" rev-parse HEAD)"
+
+      # --- choose URL to write into fetchGit ---
+      if [ "$MODE_FORCE" = "local" ]; then
+        FETCH_URL="file://$CUR_REPO_ROOT"
+      elif [ "$MODE_FORCE" = "remote" ]; then
+        if [ "$DO_PUSH" -eq 1 ]; then
+          FETCH_URL="$PUSH_REMOTE_URL"
+        elif [ "$ENTRY_EXISTS" -eq 1 ] && [ -n "$EXISTING_URL" ] && is_remote_url "$EXISTING_URL"; then
+          FETCH_URL="$EXISTING_URL"
+        else
+          CUR_ORIGIN="$(git remote get-url origin 2>/dev/null || true)"
+          [ -n "$CUR_ORIGIN" ] && is_remote_url "$CUR_ORIGIN" || die "--make-remote requires a remote url (set origin or use -p remote=URL)"
+          FETCH_URL="$CUR_ORIGIN"
+        fi
+      else
+        if [ "$DO_PUSH" -eq 1 ]; then FETCH_URL="$PUSH_REMOTE_URL"; else FETCH_URL="file://$CUR_REPO_ROOT"; fi
+      fi
+
+      # --- rewrite users.nix or inventory.nix ---
+      python3 - "$MODE" "$FILE" "$FETCH_URL" "$CUR_REV" "$USERNAME" "$DEVTYPE" "$HOSTKEY" <<'PY'
+      import sys, re, pathlib
+
+      mode = sys.argv[1]
+      path = pathlib.Path(sys.argv[2])
+      fetch_url = sys.argv[3]
+      rev = sys.argv[4]
+      username = sys.argv[5]
+      devtype = sys.argv[6]
+      hostkey = sys.argv[7]
+      text = path.read_text()
+
+      def find_matching_brace(s: str, start: int) -> int:
+          depth = 0
+          i = start
+          in_str = False
+          while i < len(s):
+              ch = s[i]
+              if in_str:
+                  if ch == '\\':
+                      i += 2
+                      continue
+                  if ch == '"':
+                      in_str = False
+                  i += 1
+                  continue
+              if ch == '"':
+                  in_str = True
+                  i += 1
+                  continue
+              if ch == '{':
+                  depth += 1
+              elif ch == '}':
+                  depth -= 1
+                  if depth == 0:
+                      return i
+              i += 1
+          raise ValueError("Could not find matching '}'")
+
+      def mk_fetch(entry_indent: str) -> str:
+        # entry_indent is indentation for the whole `"key" = <here>;` line.
+        # The attrset contents should be indented one level deeper.
+        inner = entry_indent + "  "
+        return (
+            'builtins.fetchGit {\n'
+            f'{inner}url = "{fetch_url}";\n'
+            f'{inner}rev = "{rev}";\n'
+            f'{inner}submodules = true;\n'
+            f'{entry_indent}}}'
+        )
+
+      def full_hostname(devtype: str, hostkey: str) -> str:
+          if hostkey.startswith(devtype) or "-" in hostkey:
+              return hostkey
+          if hostkey.isdigit():
+              return f"{devtype}{hostkey}"
+          return f"{devtype}-{hostkey}"
+
+      def update_user(t: str) -> str:
+          mblock = re.search(r"(?s)athenix\.users\s*=\s*\{(.*?)\n\s*\};", t)
+          if not mblock:
+              raise SystemExit("error: could not locate `athenix.users = { ... };` block")
+
+          # locate the full span of the users block to edit inside it
+          # (re-find with groups for reconstruction)
+          m2 = re.search(r"(?s)(athenix\.users\s*=\s*\{)(.*?)(\n\s*\};)", t)
+          head, body, tail = m2.group(1), m2.group(2), m2.group(3)
+
+          entry_re = re.search(
+              r"(?s)(\n[ \t]*" + re.escape(username) + r"\.external\s*=\s*)builtins\.fetchGit\s*\{",
+              body
+          )
+          if entry_re:
+              brace = body.rfind("{", 0, entry_re.end())
+              end = find_matching_brace(body, brace)
+              semi = re.match(r"\s*;", body[end+1:])
+              if not semi:
+                  raise SystemExit("error: expected ';' after fetchGit attrset")
+              semi_end = end + 1 + semi.end()
+
+              line_start = body.rfind("\n", 0, entry_re.start()) + 1
+              indent = re.match(r"[ \t]*", body[line_start:entry_re.start()]).group(0)
+
+              new_body = body[:entry_re.start()] + entry_re.group(1) + mk_fetch(indent) + ";" + body[semi_end:]
+          else:
+              indent = "    "
+              new_body = body + f"\n{indent}{username}.external = {mk_fetch(indent)};\n"
+
+          return t[:m2.start()] + head + new_body + tail + t[m2.end():]
+
+      def update_system(t: str) -> str:
+          # Find devtype block robustly: start-of-file or newline.
+          m = re.search(r"(?s)(^|\n)[ \t]*" + re.escape(devtype) + r"\s*=\s*\{", t)
+          if not m:
+              raise SystemExit(f"error: could not locate `{devtype} = {{ ... }};` block")
+
+          dev_open = t.find("{", m.end() - 1)
+          dev_close = find_matching_brace(t, dev_open)
+          dev = t[dev_open:dev_close+1]
+
+          # Find devices attrset inside dev
+          dm = re.search(r"(?s)(^|\n)[ \t]*devices\s*=\s*\{", dev)
+          if not dm:
+              raise SystemExit(f"error: could not locate `devices = {{ ... }};` inside `{devtype}`")
+
+          devices_open = dev.find("{", dm.end() - 1)
+          devices_close = find_matching_brace(dev, devices_open)
+          devices = dev[devices_open:devices_close+1]
+
+          # indentation for entries in devices
+          # find indent of the 'devices' line, then add 2 spaces
+
+          candidates = [hostkey, full_hostname(devtype, hostkey)]
+          seen = set()
+          candidates = [c for c in candidates if not (c in seen or seen.add(c))]
+
+          for key in candidates:
+              entry = re.search(
+                r'(?s)\n([ ]*)"' + re.escape(key) + r'"\s*=\s*builtins\.fetchGit\s*\{',
+                devices
+              )
+              if entry:
+                entry_indent = entry.group(1)
+
+                # find the '{' we matched
+                brace = devices.find("{", entry.end() - 1)
+                end = find_matching_brace(devices, brace)
+
+                semi = re.match(r"\s*;", devices[end+1:])
+                if not semi:
+                    raise SystemExit("error: expected ';' after fetchGit attrset in devices")
+                semi_end = end + 1 + semi.end()
+
+                # Reconstruct the prefix: newline + indent + "key" = 
+                prefix = f'\n{entry_indent}"{key}" = '
+
+                new_devices = (
+                    devices[:entry.start()]
+                    + prefix
+                    + mk_fetch(entry_indent)
+                    + ";"
+                    + devices[semi_end:]
+                )
+                new_dev = dev[:devices_open] + new_devices + dev[devices_close+1:]
+
+                return t[:dev_open] + new_dev + t[dev_close+1:]
+
+          # Not found: append into devices (exact hostkey)
+          # Indent for new entries: take indent of the closing '}' of devices, add 2 spaces.
+          close_line_start = devices.rfind("\n", 0, len(devices)-1) + 1
+          close_indent = re.match(r"[ ]*", devices[close_line_start:]).group(0)
+          entry_indent = close_indent + "  "
+
+          insertion = f'\n{entry_indent}"{hostkey}" = {mk_fetch(entry_indent)};\n'
+          new_devices = devices[:-1].rstrip() + insertion + close_indent + "}"
+          new_dev = dev[:devices_open] + new_devices + dev[devices_close+1:]
+          return t[:dev_open] + new_dev + t[dev_close+1:]
+
+      if mode == "user":
+          out = update_user(text)
+      elif mode == "system":
+          out = update_system(text)
+      else:
+          raise SystemExit("error: unknown mode")
+
+      path.write_text(out)
+      PY
+
+      cd $ATHENIX_DIR
+      nix fmt **/*.nix
+      cd $CUR_REPO_ROOT
+
+      printf "updated %s\n" "$FILE" >&2
+      printf "  url = %s\n" "$FETCH_URL" >&2
+      printf "  rev = %s\n" "$CUR_REV" >&2
+    '')
+  ];
+}

sw/updater.nix

@@ -35,25 +35,164 @@ with lib;
       (pkgs.writeShellScriptBin "update-system" ''
         set -euo pipefail
 
+        RED='\033[31m'; NC='\033[0m'
+
+        is_root() { [ "''${EUID:-$(id -u)}" -eq 0 ]; }
+        in_wheel() { id -nG 2>/dev/null | tr ' ' '\n' | grep -qx wheel; }
+
+        # Service path for unprivileged (no flags)
         UNIT="update-system.service"
 
-        # Start following logs in the background
-        journalctl -fu "$UNIT" -n 0 --output=cat &
-        JPID=$!
-
-        # Start the service and wait for it to finish
-        if systemctl start --wait --no-ask-password "$UNIT"; then
-          STATUS=$?
-        else
-          STATUS=$?
+        # Figure out the "real" invoking user, even under sudo.
+        INVOKER_USER="''${SUDO_USER:-$(id -un)}"
+        INVOKER_HOME="$(getent passwd "$INVOKER_USER" | cut -d: -f6)"
+        if [ -z "$INVOKER_HOME" ]; then
+          # fallback if getent is weird in some containers
+          INVOKER_HOME="''${HOME:-/home/$INVOKER_USER}"
         fi
 
-        sleep 2
+        # Defaults for flagged mode
+        DEFAULT_REMOTE_URL="https://git.factory.uga.edu/UGA-Innovation-Factory/athenix"
+        REPO_MODE="default"     # default | local | remote
+        LOCAL_PATH=""
+        REMOTE_URL=""
+        BRANCH=""
+        IMPURE=0
 
-        # Kill the log follower
-        kill "$JPID" 2>/dev/null || true
+        usage() {
+          cat >&2 <<'EOF'
+        usage:
+          update-system
+          update-system [--local-repo[=PATH]] [--remote-repo=URL] [--branch=BRANCH] [--impure]
 
-        exit "$STATUS"
+        notes:
+          - No flags: runs the systemd service (works for unprivileged users via polkit).
+          - Any flags: only allowed for root or wheel (runs nixos-rebuild directly).
+        EOF
+          exit 2
+        }
+
+        # No flags -> polkit-friendly systemd service route
+        if [ "$#" -eq 0 ]; then
+          journalctl -fu "$UNIT" -n 0 --output=cat &
+          JPID=$!
+
+          if systemctl start --wait --no-ask-password "$UNIT"; then
+            STATUS=$?
+          else
+            STATUS=$?
+          fi
+
+          sleep 2
+          kill "$JPID" 2>/dev/null || true
+          exit "$STATUS"
+        fi
+
+        # Flags -> require root or wheel
+        if ! is_root && ! in_wheel; then
+          printf "''${RED}error:''${NC} flags are only allowed for root or wheel. Run without flags (service path), or use sudo / add yourself to wheel.\n" >&2
+          exit 2
+        fi
+
+        # Parse flags
+        while [ "$#" -gt 0 ]; do
+          case "$1" in
+            --local-repo)
+              REPO_MODE="local"
+              LOCAL_PATH="$INVOKER_HOME/athenix"
+              shift
+              ;;
+            --local-repo=*)
+              REPO_MODE="local"
+              LOCAL_PATH="''${1#*=}"
+              shift
+              ;;
+            --remote-repo=*)
+              REPO_MODE="remote"
+              REMOTE_URL="''${1#*=}"
+              shift
+              ;;
+            --branch)
+              [ "$#" -ge 2 ] || usage
+              BRANCH="$2"
+              shift 2
+              ;;
+            --branch=*)
+              BRANCH="''${1#*=}"
+              shift
+              ;;
+            --impure)
+              IMPURE=1
+              shift
+              ;;
+            -h|--help) usage ;;
+            *)
+              printf "''${RED}error:''${NC} unknown argument: %s\n" "$1" >&2
+              usage
+              ;;
+          esac
+        done
+
+        if [ "$REPO_MODE" = "local" ] && [ -n "$REMOTE_URL" ]; then
+          printf "''${RED}error:''${NC} can't use --local-repo and --remote-repo together.\n" >&2
+          exit 2
+        fi
+
+        host="''${HOSTNAME:-$(hostname)}"
+
+        # Build flake ref
+        if [ "$REPO_MODE" = "local" ]; then
+          [ -n "$LOCAL_PATH" ] || LOCAL_PATH="$INVOKER_HOME/athenix"
+
+          # Clone default repo if missing
+          if [ ! -d "$LOCAL_PATH" ]; then
+            printf "local repo not found at %s, cloning %s...\n" "$LOCAL_PATH" "$DEFAULT_REMOTE_URL" >&2
+            if [ -n "$BRANCH" ]; then
+              git clone --branch "$BRANCH" "$DEFAULT_REMOTE_URL" "$LOCAL_PATH"
+            else
+              git clone "$DEFAULT_REMOTE_URL" "$LOCAL_PATH"
+            fi
+          fi
+
+          flakeRef="''${LOCAL_PATH}#''${host}"
+        else
+          url="''${REMOTE_URL:-$DEFAULT_REMOTE_URL}"
+
+          if echo "$url" | grep -qE '^(https?|ssh)://'; then
+            base="git+''${url}"
+          elif echo "$url" | grep -qE '^[^/@:]+@[^/:]+:'; then
+            # scp-style: git@host:owner/repo(.git)
+            userhost="''${url%%:*}"
+            path="''${url#*:}"
+            base="git+ssh://''${userhost}/''${path}"
+          else
+            base="''${url}"
+          fi
+
+          if [ -n "$BRANCH" ]; then
+            if echo "$base" | grep -q '?'; then
+              base="''${base}&ref=''${BRANCH}"
+            else
+              base="''${base}?ref=''${BRANCH}"
+            fi
+          fi
+
+          flakeRef="''${base}#''${host}"
+        fi
+
+        impureFlag=""
+        if [ "$IMPURE" -eq 1 ]; then
+          impureFlag="--impure"
+        fi
+
+        # If not root, re-exec via sudo to do the actual switch.
+        # Preserve our computed invoker context so sudo doesn't "helpfully" change it.
+        if ! is_root; then
+          exec sudo --preserve-env=HOME,USER,LOGNAME \
+            nixos-rebuild switch --refresh --print-build-logs $impureFlag --flake "$flakeRef"
+        else
+          exec nixos-rebuild switch --refresh --print-build-logs $impureFlag --flake "$flakeRef"
+        fi
       '')
     ];