#!/bin/zsh

set -euo pipefail

PKCS11_LIB="${PKCS11_LIB:-/usr/local/lib/opensc-pkcs11.so}"
YUBIKEY_PIN_FILE="${YUBIKEY_PIN_FILE:-/tmp/pass}"
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}"
OCI_PROFILE_NAME="${OCI_PROFILE_NAME:-DEFAULT}"
OCI_SESSION_VALIDATE_TIMEOUT_SECONDS="${OCI_SESSION_VALIDATE_TIMEOUT_SECONDS:-2}"
RESET_AGENT="${RESET_AGENT:-0}"
CODEX_DEVOPS_AUTH_ENV_OUT="${CODEX_DEVOPS_AUTH_ENV_OUT:-}"
DEDICATED_AGENT_PID=""
DEDICATED_AGENT_SOCK=""
PRESERVE_DEDICATED_AGENT="0"

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

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

run_oci() {
  "${OCI_BIN}" --profile "${OCI_PROFILE_NAME}" "$@"
}

resolve_timeout_bin() {
  local candidate

  for candidate in timeout gtimeout /opt/homebrew/bin/timeout /opt/homebrew/bin/gtimeout; do
    if [[ "${candidate}" == /* ]]; then
      if [[ -x "${candidate}" ]]; then
        print -r -- "${candidate}"
        return 0
      fi
      continue
    fi

    if command -v "${candidate}" >/dev/null 2>&1; then
      command -v "${candidate}"
      return 0
    fi
  done

  return 1
}

get_validate_timeout_seconds() {
  local timeout_seconds="${OCI_SESSION_VALIDATE_TIMEOUT_SECONDS}"

  if [[ ! "${timeout_seconds}" =~ '^[0-9]+([.][0-9]+)?$' ]]; then
    log "Warning: invalid OCI_SESSION_VALIDATE_TIMEOUT_SECONDS=${timeout_seconds}; using 2 seconds."
    print -r -- "2"
    return 0
  fi

  print -r -- "${timeout_seconds}"
}

run_oci_with_timeout() {
  local timeout_seconds="$1"
  shift

  local timeout_bin
  if ! timeout_bin="$(resolve_timeout_bin)"; then
    log "Warning: no timeout binary found; running OCI command without a timeout."
    run_oci "$@"
    return $?
  fi

  "${timeout_bin}" "${timeout_seconds}" "${OCI_BIN}" --profile "${OCI_PROFILE_NAME}" "$@"
}

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

  local validate_timeout_seconds
  validate_timeout_seconds="$(get_validate_timeout_seconds)"

  set +e
  run_oci_with_timeout "${validate_timeout_seconds}" 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

  if [[ ${validate_rc} -eq 124 ]]; then
    log "OCI CLI session validation timed out after ${validate_timeout_seconds} seconds; treating session as invalid."
  fi

  log "OCI CLI session is not valid; attempting refresh."
  set +e
  run_oci 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} with profile ${OCI_PROFILE_NAME}."
  "${OCI_BIN}" session authenticate --region "${OCI_SESSION_REGION}" --profile-name "${OCI_PROFILE_NAME}"
}

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

resolve_token_ssh_config_value() {
  local key="$1"

  ssh -G -F "${SSH_CONFIG_FILE}" "${TOKEN_HOST}" 2>/dev/null | awk -v key="${key}" '
    $1 == key {
      $1 = ""
      sub(/^[[:space:]]+/, "")
      print
      exit
    }
  '
}

resolve_token_proxy_command() {
  resolve_token_ssh_config_value proxycommand
}

resolve_token_proxy_jump() {
  resolve_token_ssh_config_value proxyjump
}

resolve_token_ssh_user() {
  resolve_token_ssh_config_value user
}

quiet_proxy_command() {
  local proxy_command="$1"

  proxy_command="${proxy_command// -vvv/}"
  proxy_command="${proxy_command// -vv/}"
  proxy_command="${proxy_command// -v/}"
  print -r -- "${proxy_command}"
}

proxy_ssh_prefix() {
  local ssh_bin="$1"
  local prefix="${ssh_bin} -F ${(q)SSH_CONFIG_FILE} -o LogLevel=ERROR"

  if [[ -n "${TOKEN_SSH_USER:-}" ]]; then
    prefix+=" -l ${(q)TOKEN_SSH_USER}"
  fi

  print -r -- "${prefix}"
}

proxy_command_with_ssh_config() {
  local proxy_command="$1"
  local bastion_host

  if [[ -z "${proxy_command}" || "${proxy_command}" == "none" ]]; then
    return 1
  fi

  if [[ "${proxy_command}" == bash\ -c\ * && "${proxy_command}" == *" -W %h:%p "* ]]; then
    bastion_host="${proxy_command#* -W %h:%p }"
    bastion_host="${bastion_host%% *}"
    bastion_host="${bastion_host%%\'*}"
    bastion_host="${bastion_host%%\"*}"

    if [[ -n "${bastion_host}" ]]; then
      print -r -- "$(proxy_ssh_prefix ssh) -W %h:%p ${bastion_host}"
      return 0
    fi
  fi

  case "${proxy_command}" in
    /usr/bin/ssh\ *)
      print -r -- "$(proxy_ssh_prefix /usr/bin/ssh) ${proxy_command#/usr/bin/ssh }"
      ;;
    ssh\ *)
      print -r -- "$(proxy_ssh_prefix ssh) ${proxy_command#ssh }"
      ;;
    *)
      print -r -- "${proxy_command}"
      ;;
  esac
}

proxy_jump_as_proxy_command() {
  local proxy_jump="$1"

  if [[ -z "${proxy_jump}" || "${proxy_jump}" == "none" ]]; then
    return 1
  fi

  print -r -- "$(proxy_ssh_prefix ssh) -W %h:%p ${proxy_jump}"
}

build_token_ssh_args() {
  TOKEN_SSH_ARGS=(-F "${SSH_CONFIG_FILE}" -o LogLevel=ERROR)
  TOKEN_SSH_USER="$(resolve_token_ssh_user || true)"

  local proxy_command proxy_jump configured_proxy_command
  proxy_command="$(resolve_token_proxy_command || true)"
  configured_proxy_command="$(proxy_command_with_ssh_config "${proxy_command}" || true)"

  if [[ -z "${configured_proxy_command}" ]]; then
    proxy_jump="$(resolve_token_proxy_jump || true)"
    configured_proxy_command="$(proxy_jump_as_proxy_command "${proxy_jump}" || true)"
  fi

  if [[ -n "${configured_proxy_command}" ]]; then
    TOKEN_SSH_ARGS+=(-o "ProxyCommand=${configured_proxy_command}")
  fi
}

load_pkcs11_provider_with_expect() {
  local expect_script="$1"
  shift

  if ! command -v expect >/dev/null 2>&1; then
    return 127
  fi

  if ! PKCS11_LIB_FOR_EXPECT="${PKCS11_LIB}" expect -c "${expect_script}"; then
    return 1
  fi

  return 0
}

load_pkcs11_provider_from_pin_file() {
  YUBIKEY_PIN_FILE_FOR_EXPECT="${YUBIKEY_PIN_FILE}" load_pkcs11_provider_with_expect '
        set timeout 30
        set pkcs11_lib $env(PKCS11_LIB_FOR_EXPECT)
        set pin_file $env(YUBIKEY_PIN_FILE_FOR_EXPECT)
        set fh [open $pin_file r]
        gets $fh pin
        close $fh

        spawn ssh-add -s $pkcs11_lib
        expect {
          -re {([Pp]assphrase|PIN|pin).*:} {
            send -- "$pin\r"
            exp_continue
          }
          eof {
            catch wait result
            exit [lindex $result 3]
          }
          timeout {
            exit 124
          }
        }
      ' >/dev/null
}

load_pkcs11_provider_with_prompt() {
  load_pkcs11_provider_with_expect '
        set timeout 120
        set pkcs11_lib $env(PKCS11_LIB_FOR_EXPECT)

        send_user "YubiKey PIN: "
        stty -echo
        expect_user -re "(.*)\n"
        stty echo
        send_user "\n"
        set pin $expect_out(1,string)

        log_user 0
        spawn ssh-add -s $pkcs11_lib
        expect {
          -re {([Pp]assphrase|PIN|pin).*:} {
            send -- "$pin\r"
            exp_continue
          }
          eof {
            catch wait result
            exit [lindex $result 3]
          }
          timeout {
            exit 124
          }
        }
      '
}

add_pkcs11_provider() {
  log "Loading PKCS#11 provider: ${PKCS11_LIB}"

  if [[ -r "${YUBIKEY_PIN_FILE}" ]]; then
    if ! load_pkcs11_provider_from_pin_file; then
      log "Failed to load PKCS#11 provider using YubiKey PIN file ${YUBIKEY_PIN_FILE}."
      return 1
    fi
  elif [[ -t 0 ]]; then
    if ! load_pkcs11_provider_with_prompt; then
      log "Failed to load PKCS#11 provider using prompted YubiKey PIN."
      return 1
    fi
  else
    log "YubiKey PIN file not readable at ${YUBIKEY_PIN_FILE}, and stdin is not a terminal."
    return 1
  fi

  if ! ssh-add -l >/dev/null 2>&1; then
    log "PKCS#11 provider loaded, but no SSH identities are visible to the dedicated agent."
    return 1
  fi
}

prepare_agent() {
  ensure_ssh_agent
  add_pkcs11_provider
}

refresh_operator_token() {
  local operator_token

  log "Refreshing OPERATOR_ACCESS_TOKEN from ${TOKEN_HOST} using ${SSH_CONFIG_FILE}"
  TOKEN_SSH_ARGS=()
  build_token_ssh_args

  if ! operator_token="$(ssh "${TOKEN_SSH_ARGS[@]}" "${TOKEN_HOST}" "generate --mode jwt")"; then
    log "Failed to refresh OPERATOR_ACCESS_TOKEN from ${TOKEN_HOST}."
    return 1
  fi

  if [[ -z "${operator_token}" ]]; then
    log "Token host ${TOKEN_HOST} returned an empty OPERATOR_ACCESS_TOKEN."
    return 1
  fi

  export OPERATOR_ACCESS_TOKEN="${operator_token}"
  export OP_TOKEN="${OPERATOR_ACCESS_TOKEN}"

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

write_shell_export() {
  local name="$1"
  local value="$2"

  printf 'export %s=%q\n' "${name}" "${value}"
}

write_auth_env() {
  local env_out="$1"
  local env_dir tmp

  env_dir="$(dirname -- "${env_out}")"
  mkdir -p "${env_dir}"
  tmp="$(mktemp "${env_out}.XXXXXX")"

  {
    write_shell_export SSH_AUTH_SOCK "${DEDICATED_AGENT_SOCK}"
    write_shell_export SSH_AGENT_PID "${DEDICATED_AGENT_PID}"
    write_shell_export OPERATOR_ACCESS_TOKEN "${OPERATOR_ACCESS_TOKEN}"
    write_shell_export OP_TOKEN "${OP_TOKEN}"
  } > "${tmp}"

  chmod 600 "${tmp}"
  mv -f "${tmp}" "${env_out}"
}

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

refresh_operator_token

if [[ -n "${CODEX_DEVOPS_AUTH_ENV_OUT}" ]]; then
  write_auth_env "${CODEX_DEVOPS_AUTH_ENV_OUT}"
  PRESERVE_DEDICATED_AGENT="1"
  log "Wrote Codex auth environment to ${CODEX_DEVOPS_AUTH_ENV_OUT}."
  exit 0
fi

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