#!/bin/zsh
# velcro — keep network shares stuck to your Mac.
# Remounts SMB and NFS shares after Wi-Fi drops, sleep, or network changes. No popups.
emulate -L zsh
setopt pipe_fail
zmodload zsh/system

VERSION=0.3.0
SELF=${0:A}
LABEL=io.github.dgitco.velcro
CONF_DIR=${XDG_CONFIG_HOME:-$HOME/.config}/velcro
CONF=$CONF_DIR/mounts
PAUSE=$CONF_DIR/paused
LOCK=$CONF_DIR/.lock
LOG=$HOME/Library/Logs/velcro.log
PLIST=$HOME/Library/LaunchAgents/$LABEL.plist
BIN=$HOME/.local/bin/velcro
INTERVAL=60
APP_EXEC=velcro.app/Contents/MacOS/
# Field separator for scan output. Not a tab: read collapses repeated whitespace, losing empty fields.
SEP=$'\x1f'

log() { print -r -- "$(date '+%F %T') $*" >>"$LOG"; }

# Config lines: <smb or nfs url> [fallback host ...]
# e.g.  smb://me@nas.tailnet.ts.net/home
#       smb://me@192.168.1.10/media nas.tailnet.ts.net
#       nfs://nas.local/volume1/photos
entries() {
  [[ -r $CONF ]] || return 0
  grep -vE '^[[:space:]]*(#|$)' "$CONF"
}

# Split "smb://user@host/share" or "nfs://host/export/path" into globals:
# proto user host share (URL-encoded, as in the config) name (what Finder shows).
parse() {
  proto=${1%%://*} user= host= share= name=
  [[ $proto == smb || $proto == nfs ]] || return 1
  local url=${1#*://}
  if [[ $url == *@* ]]; then user=${url%%@*}; url=${url#*@}; fi
  host=${url%%/*}
  share=${url#*/}
  [[ -n $host && -n $share && $share != "$url" ]] || return 1
  if [[ $proto == nfs ]]; then name=${share:t}; else name=$share; fi
  name=${name//\%20/ }
}

fstype() { [[ $1 == nfs ]] && print nfs || print smbfs; }
port()   { [[ $1 == nfs ]] && print 2049 || print 445; }

# Where is <share> mounted from any of <hosts...>? Prints the mount point.
# Reads `mount` output only; never touches the share, so a dead network can't hang it.
# smbfs lines look like "//me@host/My%20Share on /Volumes/My Share (smbfs, ...)",
# nfs lines like "host:/export/My Share on /Volumes/My Share (nfs, ...)".
mountpoint_of() {
  local proto=$1 share=$2 fs=$(fstype $1) h line; shift 2
  local -a pats
  for h in "$@"; do
    if [[ $proto == nfs ]]; then pats=(-e "$h:/${share//\%20/ } on ")
    else pats=(-e "@$h/$share on " -e "//$h/$share on "); fi
    line=$(mount -t $fs | grep -iF $pats | head -1)
    if [[ -n $line ]]; then
      line=${line#* on }
      print -r -- "${line% \($fs,*}"
      return 0
    fi
  done
  return 1
}

reachable() { nc -z -G 2 "$2" $(port $1) </dev/null >/dev/null 2>&1; }

# One line per share: state, share name, mount point, first host, config line (SEP-separated).
# state: mounted | reachable | offline. With "fast", unmounted shares are "offline" without probing.
scan() {
  local fast=$1 line state mp h
  local -a words hosts
  entries | while IFS= read -r line; do
    words=(${(z)line})
    parse "${words[1]}" || continue
    hosts=($host ${words[2,-1]})
    if mp=$(mountpoint_of $proto "$share" $hosts); then
      state=mounted
    else
      state=offline mp=
      if [[ -z $fast ]]; then
        for h in $hosts; do reachable $proto "$h" && { state=reachable; break; }; done
      fi
    fi
    print -r -- "$state$SEP$name$SEP$mp$SEP$host$SEP$line"
  done
}

mount_share() {
  local url=$1
  # Errors go to the log, never to a dialog. A login prompt only appears if the keychain has no password.
  osascript -e 'on run argv' -e 'try' -e 'mount volume (item 1 of argv)' \
    -e 'on error m' -e 'log m' -e 'end try' -e 'end run' "$url" </dev/null >/dev/null 2>>"$LOG"
}

cmd_run() {
  [[ -e $PAUSE || ! -r $CONF ]] && return 0
  # One pass at a time: velcro.app, the agent, and a terminal can all ask at once.
  local lockfd
  : >>"$LOCK"
  zsystem flock -t 0 -f lockfd "$LOCK" 2>/dev/null || return 0
  # Don't stack login prompts if one is already waiting.
  pgrep -f "osascript.*mount volume" >/dev/null && return 0
  [[ -f $LOG && $(stat -f %z "$LOG") -gt 1048576 ]] && mv -f "$LOG" "$LOG.1"

  local line h
  local -a words hosts
  entries | while IFS= read -r line; do
    words=(${(z)line})
    parse "${words[1]}" || { log "bad entry: $line"; continue; }
    hosts=($host ${words[2,-1]})
    mountpoint_of $proto "$share" $hosts >/dev/null && continue
    for h in $hosts; do
      reachable $proto "$h" || continue
      log "mount $proto://$h/$share"
      mount_share "$proto://${user:+$user@}$h/$share"
      break
    done
  done
  return 0
}

cmd_status() {
  local st name mp h line state agent=
  pgrep -qf "$APP_EXEC" && agent="velcro.app (on network change, wake, and a timer)"
  launchctl print "gui/$UID/$LABEL" >/dev/null 2>&1 &&
    agent+="${agent:+ + }LaunchAgent (every ${INTERVAL}s + on network change)"
  print "agent:  ${agent:-not running  →  open velcro.app, or velcro install}"
  [[ -e $PAUSE ]] && print "state:  paused  →  velcro resume" || print "state:  active"
  print
  # Grouped by server: each line in the config is one shared folder on that server.
  local -a rows; local row srv
  rows=("${(@f)$(scan)}"); [[ -z ${rows[1]} ]] && rows=()
  for srv in ${(u)${rows[@]#*$SEP*$SEP*$SEP}%%$SEP*}; do
    print -r -- "$srv"
    for row in $rows; do
      IFS=$SEP read -r st name mp h line <<<"$row"
      [[ $h == "$srv" ]] || continue
      case $st in
        mounted)   state="● mounted" ;;
        reachable) state="◐ reachable, not mounted" ;;
        *)         state="○ offline" ;;
      esac
      printf '  %-26s %-22s %s\n' "$state" "$name" "$mp"
    done
  done
  [[ -s $CONF ]] || print "no shares yet  →  velcro add smb://user@host/share"
}

cmd_add() {
  [[ -n $1 ]] || { print -u2 "usage: velcro add smb://user@host/share | nfs://host/path [fallback-host ...]"; return 2; }
  local url=${1// /%20} line
  parse "$url" || { print -u2 "not an smb:// or nfs:// url: $1"; return 2; }
  mkdir -p "$CONF_DIR"
  entries | while IFS= read -r line; do
    [[ ${${(z)line}[1]} == "$url" ]] && { print "already added: $url"; return 0; }
  done
  print -r -- "$url${2:+ ${@[2,-1]}}" >>"$CONF"
  print "added: $url${2:+ (fallback: ${@[2,-1]})}"
  cmd_run
}

# Drop the share with this exact url, or with this name if only one share has it.
cmd_rm() {
  [[ -n $1 && -r $CONF ]] || { print -u2 "usage: velcro rm <url or share name>"; return 2; }
  local arg=${1// /%20} line
  local -a words kept gone
  while IFS= read -r line || [[ -n $line ]]; do
    words=(${(z)line})
    if [[ ! $line =~ '^[[:space:]]*(#|$)' ]] && parse "${words[1]}" &&
       [[ ${words[1]} == "$arg" || $name == "$1" ]]; then
      gone+=("${words[1]}")
    else
      kept+=("$line")
    fi
  done <"$CONF"
  (( ${#gone} )) || { print -u2 "no match: $1"; return 1; }
  gone=(${(u)gone})
  if (( ${#gone} > 1 )); then
    print -u2 "more than one share is called \"$1\"; use its url:"
    print -rlu2 -- "  "${^gone}
    return 1
  fi
  if (( ${#kept} )); then print -rl -- $kept >"$CONF.tmp"; else : >"$CONF.tmp"; fi
  mv -f "$CONF.tmp" "$CONF"
  print "removed: ${gone[1]}  (still mounted? eject it in Finder)"
}

# True when this copy of velcro lives inside velcro.app.
in_app() { [[ $SELF == *.app/Contents/Resources/velcro ]]; }

cmd_install() {
  mkdir -p "${BIN:h}" "${LOG:h}" "$CONF_DIR"
  remove_swiftbar
  if in_app; then
    # The app does the reconnecting, so the agent goes and the command points into the app.
    stop_agent
    ln -sfh "$SELF" "$BIN" && print "linked: $BIN → $SELF"
    print "velcro.app reconnects your shares while it runs (it can start at login)"
  else
    # Replace, don't write through: $BIN may be a link into velcro.app.
    [[ $SELF == "${BIN:A}" ]] || { rm -f "$BIN"; cp "$SELF" "$BIN"; chmod +x "$BIN"; }
    mkdir -p "${PLIST:h}"
    cat >"$PLIST" <<EOF
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
  <key>Label</key><string>$LABEL</string>
  <key>ProgramArguments</key><array><string>$BIN</string><string>run</string></array>
  <key>RunAtLoad</key><true/>
  <key>StartInterval</key><integer>$INTERVAL</integer>
  <key>WatchPaths</key><array><string>/Library/Preferences/SystemConfiguration</string></array>
  <key>StandardOutPath</key><string>$LOG</string>
  <key>StandardErrorPath</key><string>$LOG</string>
</dict>
</plist>
EOF
    launchctl bootout "gui/$UID/$LABEL" 2>/dev/null
    # bootout is asynchronous; wait until the old agent is gone
    local i; for i in {1..20}; do launchctl print "gui/$UID/$LABEL" >/dev/null 2>&1 || break; sleep 0.25; done
    launchctl bootstrap "gui/$UID" "$PLIST" && print "installed: $BIN (agent $LABEL)"
  fi
  [[ :$PATH: == *:${BIN:h}:* ]] || print "note: add ${BIN:h} to your PATH"
}

cmd_uninstall() {
  stop_agent
  remove_swiftbar
  [[ -L $BIN ]] && rm -f "$BIN"
  print "agent removed (config kept in $CONF_DIR)"
}

stop_agent() {
  launchctl bootout "gui/$UID/$LABEL" 2>/dev/null
  rm -f "$PLIST"
}

# velcro 0.2 drew its menu with a SwiftBar plugin; velcro.app replaces it.
remove_swiftbar() {
  local dir=$(defaults read com.ameba.SwiftBar PluginDirectory 2>/dev/null)
  local plugin=${dir:-$HOME/Library/Application Support/SwiftBar/Plugins}/velcro.10s.sh
  rm -rf "$HOME/.local/share/velcro/icons"
  rmdir "$HOME/.local/share/velcro" 2>/dev/null
  [[ -e $plugin ]] || return 0
  rm -f "$plugin"
  pgrep -xq SwiftBar && open -g "swiftbar://refreshallplugins" 2>/dev/null
  print "removed the old SwiftBar menu plugin"
}

cmd_pause() {
  mkdir -p "$CONF_DIR"; : >"$PAUSE"
  print "paused — eject shares in Finder; they won't come back until 'velcro resume'"
}
cmd_resume() { rm -f "$PAUSE"; print "resumed"; cmd_run; }
cmd_logs()   { [[ -f $LOG ]] && tail -n ${1:-30} "$LOG" || print "no log yet"; }

# For velcro.app: "paused" or "active", then one scan line per share.
cmd_scan() { [[ -e $PAUSE ]] && print paused || print active; scan $1; }

# For velcro.app: the files it shares with this script.
cmd_env() {
  print -rl -- "VERSION=$VERSION" "CONF=$CONF" "PAUSE=$PAUSE" "LOG=$LOG" "BIN=$BIN" "SELF=$SELF"
}

cmd_help() {
  cat <<EOF
velcro $VERSION — keep network shares stuck to your Mac

  velcro add smb://user@host/share [fallback-host ...]
  velcro add nfs://host/export/path [fallback-host ...]
  velcro rm <url|share>      stop tracking a share
  velcro status              what's mounted, reachable, paused
  velcro pause | resume      stop/start remounting (eject stays ejected)
  velcro install | uninstall background agent (every ${INTERVAL}s + network change);
                             run from velcro.app, links this command instead
  velcro run                 one pass now
  velcro logs [n]

config: $CONF
EOF
}

case ${1:-help} in
  run)         cmd_run ;;
  status|st)   cmd_status ;;
  add)         shift; cmd_add "$@" ;;
  rm|remove)   shift; cmd_rm "$@" ;;
  install)     cmd_install ;;
  uninstall)   cmd_uninstall ;;
  pause)       cmd_pause ;;
  resume)      cmd_resume ;;
  logs|log)    shift; cmd_logs "$@" ;;
  scan)        shift; cmd_scan "$@" ;;
  env)         cmd_env ;;
  version|-v|--version) print "velcro $VERSION" ;;
  help|-h|--help) cmd_help ;;
  *) print -u2 "unknown command: $1"; cmd_help; exit 2 ;;
esac
