#!/usr/bin/env zsh

set -euo pipefail

export CODEX_HOME="${CODEX_HOME:-${HOME}/.codex}"
CODEX_AUTH_COMMON_SCRIPT="${CODEX_AUTH_COMMON_SCRIPT:-${HOME}/bin/codex-auth-common.zsh}"
CODEX_BIN="${CODEX_BIN:-/opt/homebrew/bin/codex}"
HEADROOM_BIN="${HEADROOM_BIN:-headroom}"
HEADROOM_HOST="${HEADROOM_HOST:-127.0.0.1}"
HEADROOM_PORT="${HEADROOM_PORT:-8787}"
HEADROOM_LOG_FILE="${HEADROOM_LOG_FILE:-${HOME}/.headroom/logs/proxy.log}"
CODEX_SSO_DIRECT_MODEL_PROVIDER="${CODEX_SSO_DIRECT_MODEL_PROVIDER:-openai}"
CODEX_SSO_HEADROOM_PROXY_PID=""

if [[ ! -r "${CODEX_AUTH_COMMON_SCRIPT}" ]]; then
  print -u2 -- "Codex auth common helper not found or not readable: ${CODEX_AUTH_COMMON_SCRIPT}"
  exit 1
fi
source "${CODEX_AUTH_COMMON_SCRIPT}"

codex_sso_cleanup() {
  local rc=$?

  if [[ -n "${CODEX_SSO_HEADROOM_PROXY_PID}" ]]; then
    kill "${CODEX_SSO_HEADROOM_PROXY_PID}" >/dev/null 2>&1 || true
    wait "${CODEX_SSO_HEADROOM_PROXY_PID}" >/dev/null 2>&1 || true
  fi

  codex_auth_cleanup
  return "${rc}"
}

codex_sso_headroom_ready() {
  if command -v nc >/dev/null 2>&1; then
    nc -z "${HEADROOM_HOST}" "${HEADROOM_PORT}" >/dev/null 2>&1
    return $?
  fi

  if command -v lsof >/dev/null 2>&1; then
    lsof -nP -iTCP:"${HEADROOM_PORT}" -sTCP:LISTEN >/dev/null 2>&1
    return $?
  fi

  return 1
}

codex_sso_start_headroom() {
  local attempt log_dir

  if codex_sso_headroom_ready; then
    return 0
  fi

  log_dir="$(dirname -- "${HEADROOM_LOG_FILE}")"
  mkdir -p "${log_dir}"

  "${HEADROOM_BIN}" proxy --host "${HEADROOM_HOST}" --port "${HEADROOM_PORT}" >"${HEADROOM_LOG_FILE}" 2>&1 &
  CODEX_SSO_HEADROOM_PROXY_PID=$!

  for attempt in {1..120}; do
    if codex_sso_headroom_ready; then
      return 0
    fi

    if ! kill -0 "${CODEX_SSO_HEADROOM_PROXY_PID}" >/dev/null 2>&1; then
      print -u2 -- "Headroom proxy exited before becoming ready. See ${HEADROOM_LOG_FILE}."
      return 1
    fi

    sleep 0.25
  done

  print -u2 -- "Timed out waiting for Headroom proxy on ${HEADROOM_HOST}:${HEADROOM_PORT}. See ${HEADROOM_LOG_FILE}."
  return 1
}

codex_sso_toml_string() {
  local value="$1"
  value="${value//\\/\\\\}"
  value="${value//\"/\\\"}"
  print -r -- "\"${value}\""
}

trap codex_sso_cleanup EXIT INT TERM

codex_auth_preflight

unset OPENAI_API_KEY
unset CODEX_ACCESS_TOKEN

if [[ "${CODEX_SSO_NO_HEADROOM:-}" == "1" ]]; then
  "${CODEX_BIN}" --config "model_provider=$(codex_sso_toml_string "${CODEX_SSO_DIRECT_MODEL_PROVIDER}")" --disable guardian_approval -a untrusted "$@"
else
  codex_sso_start_headroom
  OPENAI_BASE_URL="http://${HEADROOM_HOST}:${HEADROOM_PORT}/v1" "${CODEX_BIN}" --disable guardian_approval -a untrusted "$@"
fi
# "${CODEX_BIN}" "$@"