# Get the filename to store/lookup the environment from
ssh_env_cache="$HOME/.ssh/environment-$SHORT_HOST"

# Test if $SSH_AUTH_SOCK points at a listening agent
function _is_agent_running() {
  local REPLY
  [[ -n "$SSH_AUTH_SOCK" && -S "$SSH_AUTH_SOCK" ]] || return 1
  zmodload zsh/net/socket 2>/dev/null || return 1
  zsocket "$SSH_AUTH_SOCK" 2>/dev/null || return 1
  exec {REPLY}>&- # close the probe connection
  return 0
}

function _start_agent() {
  # Reuse an agent the session already provides, if enabled
  if zstyle -t :omz:plugins:ssh-agent honor-existing && _is_agent_running; then
    return 0
  fi

  # Check if ssh-agent is already running
  if [[ -f "$ssh_env_cache" ]]; then
    . "$ssh_env_cache" > /dev/null

    # Test if $SSH_AUTH_SOCK is visible
    _is_agent_running && return 0
  fi

  if [[ ! -d "$HOME/.ssh" ]]; then
    echo "[oh-my-zsh] ssh-agent plugin requires ~/.ssh directory"
    return 1
  fi

  # Set a maximum lifetime for identities added to ssh-agent
  local lifetime
  zstyle -s :omz:plugins:ssh-agent lifetime lifetime

  # start ssh-agent and setup environment
  zstyle -t :omz:plugins:ssh-agent quiet || echo >&2 "Starting ssh-agent ..."
  ssh-agent -s ${lifetime:+-t} ${lifetime} | sed '/^echo/d' >! "$ssh_env_cache"
  chmod 600 "$ssh_env_cache"
  . "$ssh_env_cache" > /dev/null
}

function _add_identities() {
  local id file line sig lines
  local -a identities loaded_sigs loaded_ids not_loaded
  zstyle -a :omz:plugins:ssh-agent identities identities

  # check for .ssh folder presence
  if [[ ! -d "$HOME/.ssh" ]]; then
    return
  fi

  # add default keys if no identities were set up via zstyle
  # this is to mimic the call to ssh-add with no identities
  if [[ ${#identities} -eq 0 ]]; then
    # key list found on `ssh-add` man page's DESCRIPTION section
    for id in id_rsa id_dsa id_ecdsa id_ed25519 id_ed25519_sk identity; do
      # check if file exists
      [[ -f "$HOME/.ssh/$id" ]] && identities+=($id)
    done
  fi

  # get list of loaded identities' signatures and filenames
  if lines=$(ssh-add -l); then
    for line in ${(f)lines}; do
      loaded_sigs+=${${(z)line}[2]}
      loaded_ids+=${${(z)line}[3]}
    done
  fi

  # add identities if not already loaded
  for id in $identities; do
    # if id is an absolute path, make file equal to id
    [[ "$id" = /* ]] && file="$id" || file="$HOME/.ssh/$id"
    # check for filename match, otherwise try for signature match
    if [[ -f $file && ${loaded_ids[(I)$file]} -le 0 ]]; then
      sig="$(ssh-keygen -lf "$file" | awk '{print $2}')"
      [[ ${loaded_sigs[(I)$sig]} -le 0 ]] && not_loaded+=("$file")
    fi
  done

  # abort if no identities need to be loaded
  if [[ ${#not_loaded} -eq 0 ]]; then
    return
  fi

  # pass extra arguments to ssh-add
  local args
  zstyle -a :omz:plugins:ssh-agent ssh-add-args args

  # if ssh-agent quiet mode, pass -q to ssh-add
  zstyle -t :omz:plugins:ssh-agent quiet && args=(-q $args)

  # use user specified helper to ask for password (ksshaskpass, etc)
  local helper
  zstyle -s :omz:plugins:ssh-agent helper helper

  if [[ -n "$helper" ]]; then
    if [[ -z "${commands[$helper]}" ]]; then
      echo >&2 "ssh-agent: the helper '$helper' has not been found."
    else
      SSH_ASKPASS="$helper" ssh-add "${args[@]}" ${^not_loaded} < /dev/null
      return $?
    fi
  fi

  ssh-add "${args[@]}" ${^not_loaded}
}

# Add a nifty symlink for screen/tmux if agent forwarding is enabled
if zstyle -t :omz:plugins:ssh-agent agent-forwarding \
   && [[ -n "$SSH_AUTH_SOCK" ]]; then
  if [[ ! -L "$SSH_AUTH_SOCK" ]]; then
    if [[ -n "$TERMUX_VERSION" ]]; then
      _omz_ssh_agent_link="$PREFIX"/tmp/ssh-agent-$USERNAME-screen
    else
      _omz_ssh_agent_link=/tmp/ssh-agent-$USERNAME-screen
    fi
    # `ln -sf` unlinks the old symlink and then creates the new one: two
    # syscalls, no atomic swap. Shells starting concurrently (a tmux window
    # opening several panes at once) interleave those steps and all but one
    # fail with "ln: ...: File exists". Stage the link under a PID-unique
    # name and move it into place, since rename(2) is atomic.
    command ln -sf "$SSH_AUTH_SOCK" "$_omz_ssh_agent_link.$$" \
      && command mv -f "$_omz_ssh_agent_link.$$" "$_omz_ssh_agent_link"
    unset _omz_ssh_agent_link
  fi
else
  _start_agent
fi

# Don't add identities if lazy-loading is enabled
if ! zstyle -t :omz:plugins:ssh-agent lazy; then
  _add_identities
fi

unset agent_forwarding ssh_env_cache
unfunction _start_agent _add_identities _is_agent_running