#!/bin/zsh

set -euo pipefail

PKCS11_LIB="${PKCS11_LIB:-/usr/local/lib/opensc-pkcs11.so}"
TOKEN_HOST="${TOKEN_HOST:-operator-access-token.svc.ad1.r2}"
SSH_CONFIG_FILE="${SSH_CONFIG_FILE:-$HOME/.ssh/config.oci}"
OCI_BIN="${OCI_BIN:-/opt/homebrew/bin/oci}"
OCI_SESSION_REGION="${OCI_SESSION_REGION:-us-chicago-1}"
RESET_AGENT="${RESET_AGENT:-0}"
DEDICATED_AGENT_PID=""

log() {
  print -u2 -- "$@"
}

inherited_agent_likely_rejects_pkcs11() {
  [[ -n "${SSH_AUTH_SOCK:-}" ]] || return 1
  [[ -z "${SSH_AGENT_PID:-}" ]] || return 1

  case "${SSH_AUTH_SOCK}" in
    /private/tmp/com.apple.launchd.*/Listeners)
      return 0
      ;;
  esac

  return 1
}

cleanup() {
  if [[ -n "${DEDICATED_AGENT_PID}" ]]; then
    SSH_AGENT_PID="${DEDICATED_AGENT_PID}" ssh-agent -k >/dev/null 2>&1 || true
  fi
}

ensure_oci_session() {
  if [[ ! -x "${OCI_BIN}" ]]; then
    print -u2 "OCI CLI not found or not executable: ${OCI_BIN}"
    exit 1
  fi

  set +e
  "${OCI_BIN}" session validate >/dev/null 2>&1
  local validate_rc=$?
  set -e

  if [[ ${validate_rc} -eq 0 ]]; then
    log "OCI CLI session is already valid."
    return 0
  fi

  log "OCI CLI session is not valid; attempting refresh."
  set +e
  "${OCI_BIN}" session refresh >/dev/null 2>&1
  local refresh_rc=$?
  set -e

  if [[ ${refresh_rc} -eq 0 ]]; then
    log "OCI CLI session refresh succeeded."
    return 0
  fi

  log "Running OCI CLI session authenticate for ${OCI_SESSION_REGION}."
  "${OCI_BIN}" session authenticate --region "${OCI_SESSION_REGION}"
}

ensure_ssh_agent() {
  if [[ -n "${SSH_AUTH_SOCK:-}" && -S "${SSH_AUTH_SOCK}" ]]; then
    set +e
    ssh-add -l >/dev/null 2>&1
    local rc=$?
    set -e
    case ${rc} in
      0|1)
        return 0
        ;;
    esac
  fi

  log "Starting ssh-agent for Codex."
  eval "$(ssh-agent -s)" >/dev/null
  DEDICATED_AGENT_PID="${SSH_AGENT_PID:-}"
}

add_pkcs11_provider() {
  log "Loading PKCS#11 provider: ${PKCS11_LIB}"
  ssh-add -s "${PKCS11_LIB}" >/dev/null
}

prepare_agent() {
  local had_inherited_agent=0
  if [[ -n "${SSH_AUTH_SOCK:-}" && -S "${SSH_AUTH_SOCK}" ]]; then
    had_inherited_agent=1
  fi

  if inherited_agent_likely_rejects_pkcs11; then
    log "Inherited launchd SSH agent is unlikely to support PKCS#11; starting a dedicated ssh-agent for Codex."
    unset SSH_AUTH_SOCK SSH_AGENT_PID
    had_inherited_agent=0
  fi

  ensure_ssh_agent

  if add_pkcs11_provider; then
    return 0
  fi

  if [[ ${had_inherited_agent} -eq 1 ]]; then
    log "Existing SSH agent rejected PKCS#11 provider; starting a dedicated ssh-agent for Codex."
    unset SSH_AUTH_SOCK SSH_AGENT_PID
    ensure_ssh_agent
    add_pkcs11_provider
    return 0
  fi

  return 1
}

if [[ ! -f "${SSH_CONFIG_FILE}" ]]; then
  print -u2 "SSH config file not found: ${SSH_CONFIG_FILE}"
  exit 1
fi

if [[ "${RESET_AGENT}" == "1" ]]; then
  log "Resetting SSH agent on explicit request."
  pkill -9 ssh-agent >/dev/null 2>&1 || true
  pkill -9 ssh-pkcs11-helper >/dev/null 2>&1 || true
  sleep 1
fi

trap cleanup EXIT INT TERM

ensure_oci_session

prepare_agent

log "Refreshing OPERATOR_ACCESS_TOKEN from ${TOKEN_HOST} using ${SSH_CONFIG_FILE}"
export OPERATOR_ACCESS_TOKEN="$(ssh -F "${SSH_CONFIG_FILE}" "${TOKEN_HOST}" "generate --mode jwt")"
export OP_TOKEN="${OPERATOR_ACCESS_TOKEN}"

log "Using fresh OP_TOKEN for Codex and DevOps MCP."

/opt/homebrew/bin/codex "$@"