#!/usr/bin/env bash
# shellcheck disable=SC2034
# SPDX-License-Identifier: MPL-2.0
# Reliably detect and switch between SSH agents like keychain, yubikey-agent, 1Password, etc.

if [[ $DEBUG != "" ]]; then
  set -x
fi

logOps() {
  PREFIX=$3

  if [[ $PREFIX != "" ]]; then
    LOGOPS_PREFIX="[ssh-agent-loader::$PREFIX]"
  else
    LOGOPS_PREFIX="[ssh-agent-loader]"
  fi

  if [[ $1 == "debug" ]]; then
    [[ $DEBUG != "" ]] && echo "[ssh-agent-loader] debug: $2"
  elif [[ $1 == "warn" ]]; then
    echo "$LOGOPS_PREFIX warning: $2"
  elif [[ $1 == "error" ]]; then
    echo "$LOGOPS_PREFIX error: $2"
  else
    [[ $SSH_AGENT_LOADER_SLIENT != "1" ]] && echo "$LOGOPS_PREFIX $2"
  fi
}

# Workaround in cases where XDG_RUNTIME_DIR is undefined on login
if [[ -z "${XDG_RUNTIME_DIR}" ]]; then
  logOps warn "XDG_RUNTIME_DIR is possibly undefined, see https://github.com/swaywm/sway/issues/7202"
  logOps warn "for context and https://wiki.archlinux.org/title/XDG_Base_Directory for docs"
  logOps warn "Temporarily using '/run/user/$(id -u)' for XDG_RUNTIME_DIR in 3s..."
  sleep 3
  XDG_RUNTIME_DIR="/run/user/$(id -u)"
  export XDG_RUNTIME_DIR
fi

# do feature detection if keychain is installed
if command -v keychain >> /dev/null; then
  FF_KEYCHAIN=1
else
  FF_KEYCHAIN=0
fi

try_keychain_ssh_agent() {
  if [[ $FF_KEYCHAIN == "1" ]]; then
    logOps info "attempting to use keychain for SSH agents" keychain
    if [[ "$SSH_AGENT_LOADER_SLIENT" == "" ]]; then
      eval "$(keychain --eval --ssh-spawn-gpg --ssh-allow-gpg)"
    else
      eval "$(keychain --eval --ssh-spawn-gpg --ssh-allow-gpg --quiet)"
    fi
  else
    logOps warn "keychain is not in PATH yet" keychain
    return 1
  fi
}

# Ripped off NixOS-generated set-environment on my laptop for yubikey-agent setup
try_yubikey_agent() {
  YUBIKEY_AGENT_AUTH_SOCK="${XDG_RUNTIME_DIR}/yubikey-agent/yubikey-agent.sock"
  if [[ -f "${YUBIKEY_AGENT_AUTH_SOCK}" ]]; then
    logOps info "attempting to use Yubikey SSH agent" yubikey-agent
    if ! SSH_AUTH_SOCK=${YUBIKEY_AGENT_AUTH_SOCK} ssh-add -l >> /dev/null 2>&1; then
      logOps warn "something went wrong while checking for Yubikey SSH agent availability" yubikey-agent
      logOps warn "is your Yubikey plugged in properly?" yubikey-agent
      return 1
    fi
  else
    logOps error "Yubikey SSH agent seems to be not available on this host" yubikey-agent
    return 1
  fi
}

try_1password_ssh_agent() {
  OP_SSH_AUTH_SOCK="$HOME/.1password/agent.sock"
  if [[ ! -S "$OP_SSH_AUTH_SOCK" ]]; then
    logOps warn "1Password SSH agent isn't enabled or desktop app isn't installed yet" 1password
    return 1
  fi

  logOps info "attempting to use 1Password SSH agent" 1password
  if ! SSH_AUTH_SOCK=$OP_SSH_AUTH_SOCK ssh-add -l >> /dev/null 2>&1; then
    logOps warn "something went wrong while checking for 1Password SSH agent availability" 1password
    logOps warn "unlock the desktop app first or enable SSH agent from settings" 1password
    return 1
  fi
  export SSH_AUTH_SOCK=$OP_SSH_AUTH_SOCK FF_USE_OP_CLI_PLUGINS=true
  unset OP_SSH_AUTH_SOCK
}

ssh-agent-loader() {
  if [[ $1 == "" || $1 == "auto" ]]; then
    if [[ $SSH_CONNECTION != "" ]] && [[ $VSCODE_IPC_HOOK_CLI != "" ]]; then
      logOps info "automatic detection is disabled while you're in a VS Code Remote SSH/Tunnels session"
      logOps info "to enable SSH agent, please manually invoke the shell function with desired agent"
      return
    fi

    export OLD_SSH_AUTH_SOCK="$SSH_AUTH_SOCK"

    unset SSH_AGENT_PID SSH_AUTH_SOCK
    if try_1password_ssh_agent; then
      return
    elif try_keychain_ssh_agent; then
      return
    elif try_yubikey_agent; then
      return
    else
      logOps error "SSH agent seems to be failed to load at the moment"
      logOps error "try again later by manually invoking the shell function"
      return 1
    fi
  elif [[ $1 == "1passowrd" || $1 == "op" ]]; then
    unset SSH_AGENT_PID SSH_AUTH_SOCK
    try_1password_ssh_agent
  elif [[ $1 == "keychain" ]]; then
    try_keychain_ssh_agent
  elif [[ $1 == "yubikey" ]]; then
    try_yubikey_agent
  else
    echo "Usage: ssh-agent-loader [auto|[1password|op|1p]|keychain|yubikey]"
    return 1
  fi
}

# Avoid source loops on subsequent file resourcing.
if [[ -z "$SSH_AGENT_LOADED" ]]; then
  # automatically detect things as we source this
  [[ "$FF_SKIP_AUTO_SSH_AGENT_LOADER" == "" ]] && ssh-agent-loader auto
  export SSH_AGENT_LOADED=1
fi

if [[ $DEBUG == "1" ]]; then
  set +x
fi