sw/update-ref.nix

@@ -1,7 +1,8 @@
 { pkgs, ... }:
 {
-  environment.systemPackages = [
+  environment.systemPackages = with pkgs; [
-    pkgs.python3
+    python3
+    git
     (pkgs.writeShellScriptBin "update-ref" ''
       set -euo pipefail
 
@@ -9,46 +10,37 @@
       die() { printf "''${RED}error:''${NC} %s\n" "$*" >&2; exit 2; }
       warn() { printf "''${YEL}warning:''${NC} %s\n" "$*" >&2; }
 
-      # ---- Must be in a git repo (current dir) ----
+      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 ----
+      # --- athenix checkout (working tree) ---
       ATHENIX_DIR="$HOME/athenix"
-      ATHENIX_REPO=""
       ATHENIX_BRANCH=""
 
-      # ---- Git automation options for CURRENT repo ----
+      # --- current repo automation ---
       COMMIT_MSG=""
       PUSH_SPEC=""
 
-      # ---- Push + mode controls ----
+      # --- push / url mode ---
-      PUSH_SET=0            # user explicitly set push true/false
+      PUSH_SET=0
-      DO_PUSH=0             # final push decision
+      DO_PUSH=0
-      MODE_FORCE=""         # "", "local", "remote" (from -l/-r)
+      MODE_FORCE=""  # "", local, remote
 
-      # ---- Required subcommand ----
+      TARGET=""
-      TARGET=""             # user=<username> OR system=<device-type>:<hostkey>
-
-      usage() {
-        cat >&2 <<'EOF'
-      usage:
-        update-ref [--athenix-repo=PATH|URL|-R PATH|URL] [--athenix-branch=BRANCH|-b 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>
-
-      push default:
-        - determined from the existing builtins.fetchGit url in the target entry:
-            remote url -> push=true
-            file/local -> push=false
-        - if the entry doesn't exist: push=false unless --make-remote
-
-      EOF
-        exit 2
-      }
 
       is_remote_url() {
         # https://, http://, ssh://, or scp-style git@host:org/repo
@@ -56,8 +48,7 @@
       }
 
       derive_full_hostname() {
-        devtype="$1"
+        devtype="$1"; hostkey="$2"
-        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
@@ -71,77 +62,57 @@
         # 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()
 
-      mode = sys.argv[1]      # user | system
+      def url_from_block(block: str) -> str:
-      path = pathlib.Path(sys.argv[2])
-      username = sys.argv[3]  # user mode
-      key = sys.argv[4]       # system mode: exact string key
-
-      text = path.read_text()
-
-      def find_fetch_block_user(u):
-          m = re.search(r'(?s)\n\s*' + re.escape(u) + r'\.external\s*=\s*builtins\.fetchGit\s*\{(.*?)\n\s*\};', text)
-          return m.group(1) if m else None
-
-      def find_fetch_block_system(k):
-          m = re.search(r'(?s)\n\s*"' + re.escape(k) + r'"\s*=\s*builtins\.fetchGit\s*\{(.*?)\n\s*\};', text)
-          return m.group(1) if m else None
-
-      def extract_url(block):
           if not block:
               return ""
           m = re.search(r'url\s*=\s*"([^"]+)"\s*;', block)
           return m.group(1) if m else ""
 
-      block = None
       if mode == "user":
-          block = find_fetch_block_user(username)
+          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:
-          block = find_fetch_block_system(key)
+          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(extract_url(block))
+          print(url_from_block(block))
       PY
       }
 
-      # ---- Parse args ----
+      # --- parse args ---
       while [ "$#" -gt 0 ]; do
         case "$1" in
           user=*|system=*)
             [ -z "$TARGET" ] || die "Only one subcommand allowed (user=... or system=...)"
-            TARGET="$1"
+            TARGET="$1"; shift
-            shift
             ;;
           --athenix-repo=*)
-            ATHENIX_REPO="''${1#*=}"
+            ATHENIX_DIR="''${1#*=}"; shift
-            shift
             ;;
           -R)
             [ "$#" -ge 2 ] || usage
-            ATHENIX_REPO="$2"
+            ATHENIX_DIR="$2"; shift 2
-            shift 2
             ;;
           --athenix-branch=*)
-            ATHENIX_BRANCH="''${1#*=}"
+            ATHENIX_BRANCH="''${1#*=}"; shift
-            shift
             ;;
           -b)
             [ "$#" -ge 2 ] || usage
-            ATHENIX_BRANCH="$2"
+            ATHENIX_BRANCH="$2"; shift 2
-            shift 2
             ;;
           -m|--message)
             [ "$#" -ge 2 ] || usage
-            COMMIT_MSG="$2"
+            COMMIT_MSG="$2"; shift 2
-            shift 2
             ;;
 
           -p|--push)
-            PUSH_SET=1
+            PUSH_SET=1; DO_PUSH=1
-            DO_PUSH=1
             # optional next token remote or remote=URL (only if it doesn't look like another flag/subcommand)
-            if [ "$#" -ge 2 ] && printf "%s" "$2" | grep -qvE '^(user=|system=|--|-l$|-r$)'; then
+            if [ "$#" -ge 2 ] && printf "%s" "$2" | grep -qvE '^(user=|system=|--|-l$|-r$)$'; then
-              PUSH_SPEC="$2"
+              PUSH_SPEC="$2"; shift 2
-              shift 2
             else
               shift
             fi
@@ -157,15 +128,8 @@
             shift
             ;;
 
-          --make-local|-l)
+          --make-local|-l) MODE_FORCE="local"; shift ;;
-            MODE_FORCE="local"
+          --make-remote|-r) MODE_FORCE="remote"; shift ;;
-            shift
-            ;;
-          --make-remote|-r)
-            MODE_FORCE="remote"
-            shift
-            ;;
-
           -h|--help) usage ;;
           *) die "Unknown argument: $1" ;;
         esac
@@ -173,20 +137,35 @@
 
       [ -n "$TARGET" ] || die "Missing required subcommand: user=<username> or system=<device-type>:<hostkey>"
 
-      # ---- Validate athenix checkout ----
+      # --- 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"
+      git -C "$ATHENIX_DIR" rev-parse --is-inside-work-tree >/dev/null 2>&1 || die "$ATHENIX_DIR is not a git project (athenix checkout)"
-      if [ -z "$ATHENIX_BRANCH" ]; then
+
-        ATHENIX_BRANCH="$(git -C "$ATHENIX_DIR" rev-parse --abbrev-ref HEAD)"
+      # --- -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
 
-      # ---- Determine target file + identifiers ----
+      # --- target file + identifiers ---
-      FILE=""
+      MODE=""; FILE=""; USERNAME=""; DEVTYPE=""; HOSTKEY=""
-      MODE=""
-      USERNAME=""
-      DEVTYPE=""
-      HOSTKEY=""
-
       case "$TARGET" in
         user=*)
           MODE="user"
@@ -205,13 +184,11 @@
           FILE="$ATHENIX_DIR/inventory.nix"
           ;;
       esac
-
       [ -f "$FILE" ] || die "File not found: $FILE"
 
-      # ---- Determine existing fetchGit url in the entry (for push default) ----
+      # --- 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
@@ -226,26 +203,19 @@
         fi
       fi
 
-      # ---- Default push based on existing entry url unless user explicitly set it ----
       if [ "$PUSH_SET" -eq 0 ]; then
         if [ "$ENTRY_EXISTS" -eq 1 ] && is_remote_url "$EXISTING_URL"; then
           DO_PUSH=1
         else
-          # entry missing or local/file url
+          DO_PUSH=0
-          if [ "$MODE_FORCE" = "remote" ]; then
+          [ "$MODE_FORCE" = "remote" ] && DO_PUSH=1 || true
-            DO_PUSH=1
-          else
-            DO_PUSH=0
-          fi
         fi
       fi
-
-      # If forcing local and user didn't explicitly demand push, push defaults off.
       if [ "$MODE_FORCE" = "local" ] && [ "$PUSH_SET" -eq 0 ]; then
         DO_PUSH=0
       fi
 
-      # ---- Dirty check prompt ----
+      # --- 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
@@ -253,25 +223,16 @@
           y|Y|yes|YES)
             git add -A
             if ! git diff --cached --quiet; then
-              if [ -n "$COMMIT_MSG" ]; then
+              if [ -n "$COMMIT_MSG" ]; then git commit -m "$COMMIT_MSG"; else git commit; fi
-                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.)" ;;
-            warn "Proceeding without committing. (rev will be last committed HEAD.)"
-            ;;
         esac
-      elif [ -n "$COMMIT_MSG" ]; then
-        # clean tree but message provided: nothing to do
-        :
       fi
 
-      # ---- Push current repo if requested ----
+      # --- push current repo if requested ---
       PUSH_REMOTE_URL=""
       if [ "$DO_PUSH" -eq 1 ]; then
         if [ -n "$PUSH_SPEC" ]; then
@@ -280,13 +241,11 @@
             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
@@ -295,7 +254,6 @@
             PUSH_REMOTE_URL="$(git remote get-url "$REM_NAME")"
           fi
         else
-          # default: push to upstream, error if none
           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
@@ -305,10 +263,9 @@
         fi
       fi
 
-      # ---- Always get current HEAD hash (after optional commit/push) ----
       CUR_REV="$(git -C "$CUR_REPO_ROOT" rev-parse HEAD)"
 
-      # ---- Determine url to write into fetchGit (respecting make-local/make-remote) ----
+      # --- choose URL to write into fetchGit ---
       if [ "$MODE_FORCE" = "local" ]; then
         FETCH_URL="file://$CUR_REPO_ROOT"
       elif [ "$MODE_FORCE" = "remote" ]; then
@@ -322,14 +279,10 @@
           FETCH_URL="$CUR_ORIGIN"
         fi
       else
-        if [ "$DO_PUSH" -eq 1 ]; then
+        if [ "$DO_PUSH" -eq 1 ]; then FETCH_URL="$PUSH_REMOTE_URL"; else FETCH_URL="file://$CUR_REPO_ROOT"; fi
-          FETCH_URL="$PUSH_REMOTE_URL"
-        else
-          FETCH_URL="file://$CUR_REPO_ROOT"
-        fi
       fi
 
-      # ---- Rewrite athenix file ----
+      # --- rewrite users.nix or inventory.nix ---
       python3 - "$MODE" "$FILE" "$FETCH_URL" "$CUR_REV" "$USERNAME" "$DEVTYPE" "$HOSTKEY" <<'PY'
       import sys, re, pathlib
 
@@ -340,75 +293,169 @@
       username = sys.argv[5]
       devtype = sys.argv[6]
       hostkey = sys.argv[7]
-
       text = path.read_text()
 
-      def mk_fetch_block(url, rev):
+      def find_matching_brace(s: str, start: int) -> int:
-          return (
+          depth = 0
-              'builtins.fetchGit {\n'
+          i = start
-              f'      url = "{url}";\n'
+          in_str = False
-              f'      rev = "{rev}";\n'
+          while i < len(s):
-              '      submodules=true;\n'
+              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 '}'")
 
-      fetch_block = mk_fetch_block(fetch_url, rev)
+      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, hostkey):
+      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}"
 
-      if mode == "user":
+      def update_user(t: str) -> str:
-          m = re.search(r'(?s)(athenix\.users\s*=\s*\{)(.*?)(\n\s*\};)', text)
+          mblock = re.search(r"(?s)athenix\.users\s*=\s*\{(.*?)\n\s*\};", t)
-          if not m:
+          if not mblock:
               raise SystemExit("error: could not locate `athenix.users = { ... };` block")
 
-          head, body, tail = m.group(1), m.group(2), m.group(3)
+          # 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)
 
-          key_re = re.compile(r'(?s)(\n\s*' + re.escape(username) + r'\.external\s*=\s*)builtins\.fetchGit\s*\{.*?\};')
+          entry_re = re.search(
-          if key_re.search(body):
+              r"(?s)(\n[ \t]*" + re.escape(username) + r"\.external\s*=\s*)builtins\.fetchGit\s*\{",
-              body = key_re.sub(lambda mm: mm.group(1) + fetch_block + ';', body)
+              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:
-              body = body + f'\n    {username}.external = {fetch_block};\n'
+              indent = "    "
+              new_body = body + f"\n{indent}{username}.external = {mk_fetch(indent)};\n"
 
-          new_text = text[:m.start()] + head + body + tail + text[m.end():]
+          return t[:m2.start()] + head + new_body + tail + t[m2.end():]
-          path.write_text(new_text)
-          sys.exit(0)
 
-      elif mode == "system":
+      def update_system(t: str) -> str:
-          block_re = re.compile(r'(?s)(\n\s*' + re.escape(devtype) + r'\s*=\s*\{)(.*?)(\n\s*\};)')
+          # Find devtype block robustly: start-of-file or newline.
-          m = block_re.search(text)
+          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 in inventory.nix")
+              raise SystemExit(f"error: could not locate `{devtype} = {{ ... }};` block")
 
-          head, body, tail = m.group(1), m.group(2), m.group(3)
+          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)]
-          candidates = list(dict.fromkeys(candidates))
+          seen = set()
+          candidates = [c for c in candidates if not (c in seen or seen.add(c))]
 
-          replaced = False
+          for key in candidates:
-          for k in candidates:
+              entry = re.search(
-              entry_re = re.compile(r'(?s)(\n\s*"' + re.escape(k) + r'"\s*=\s*)builtins\.fetchGit\s*\{.*?\};')
+                r'(?s)\n([ ]*)"' + re.escape(key) + r'"\s*=\s*builtins\.fetchGit\s*\{',
-              if entry_re.search(body):
+                devices
-                  body = entry_re.sub(lambda mm: mm.group(1) + fetch_block + ';', body)
+              )
-                  replaced = True
+              if entry:
-                  break
+                entry_indent = entry.group(1)
 
-          if not replaced:
+                # find the '{' we matched
-              body = body + f'\n      "{hostkey}" = {fetch_block};\n'
+                brace = devices.find("{", entry.end() - 1)
+                end = find_matching_brace(devices, brace)
 
-          new_text = text[:m.start()] + head + body + tail + text[m.end():]
+                semi = re.match(r"\s*;", devices[end+1:])
-          path.write_text(new_text)
+                if not semi:
-          sys.exit(0)
+                    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
 
-      printf "updated %s (athenix branch: %s)\n" "$FILE" "$ATHENIX_BRANCH" >&2
+      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
     '')