#!/usr/bin/env bash
#
# gh CLI wrapper that fetches fresh GitHub tokens from agent-relay API
# Caches tokens for 55 minutes to minimize API calls.
#
# This solves the problem of expired GH_TOKEN in workspace environments.
# The provisioner sets GH_TOKEN at startup, but installation tokens expire
# after ~1 hour. This wrapper fetches fresh tokens before each gh invocation.
#
# Install: Copy to /usr/local/bin/gh (before /usr/bin/gh in PATH)
#

set -euo pipefail

CACHE_DIR="${HOME}/.cache/gh-relay"
CACHE_FILE="${CACHE_DIR}/token"
CACHE_TTL_SECONDS=3300  # 55 minutes (tokens expire at 1 hour)
AUTH_ERROR_PATTERN='(HTTP 401|401 Unauthorized|Bad credentials|authentication failed|invalid token)'

# Ensure cache directory exists
mkdir -p "$CACHE_DIR"

# Check if we have required env vars for relay token refresh
use_relay() {
  [[ -n "${WORKSPACE_ID:-}" ]] && [[ -n "${CLOUD_API_URL:-}" ]] && [[ -n "${WORKSPACE_TOKEN:-}" ]]
}

# Get cached token if still valid
get_cached_token() {
  if [[ -f "$CACHE_FILE" ]]; then
    local cached_time
    cached_time=$(stat -c %Y "$CACHE_FILE" 2>/dev/null || stat -f %m "$CACHE_FILE" 2>/dev/null || echo 0)
    local now
    now=$(date +%s)
    local age=$((now - cached_time))

    if [[ $age -lt $CACHE_TTL_SECONDS ]]; then
      cat "$CACHE_FILE"
      return 0
    fi
  fi
  return 1
}

# Fetch fresh token from API
fetch_fresh_token() {
  local response
  response=$(curl -sf \
    -H "Authorization: Bearer ${WORKSPACE_TOKEN}" \
    "${CLOUD_API_URL}/api/git/token?workspaceId=${WORKSPACE_ID}" \
    2>/dev/null) || return 1

  # Prefer userToken for gh CLI (works for user-context operations like pr create)
  # Fall back to token field (which is also userToken-first since the API change)
  local token
  token=$(echo "$response" | jq -r '.userToken // .token // empty')

  if [[ -n "$token" ]]; then
    echo "$token" > "$CACHE_FILE"
    chmod 600 "$CACHE_FILE"
    echo "$token"
    return 0
  fi
  return 1
}

# Clear cached token (e.g., when we detect an auth failure)
clear_cached_token() {
  rm -f "$CACHE_FILE"
}

# Get token (cached or fresh)
get_token() {
  # Try cache first
  local token
  if token=$(get_cached_token); then
    echo "$token"
    return 0
  fi

  # Fetch fresh token with retry (silent)
  local attempt=0
  local max_retries=5
  local delays=(1 2 4 8 16)
  while [[ $attempt -lt $max_retries ]]; do
    if token=$(fetch_fresh_token); then
      echo "$token"
      return 0
    fi
    sleep "${delays[$attempt]}"
    attempt=$((attempt + 1))
  done

  return 1
}

# Run gh and capture stderr to detect auth failures without losing output.
# Sets AUTH_ERROR=1 when auth errors are detected.
AUTH_ERROR=0
run_gh() {
  local err_file
  err_file=$(mktemp)
  AUTH_ERROR=0

  if "$GH_BIN" "$@" 2> >(tee "$err_file" >&2); then
    :
  else
    local status=$?
    if grep -qiE "$AUTH_ERROR_PATTERN" "$err_file"; then
      AUTH_ERROR=1
    fi
    rm -f "$err_file"
    return "$status"
  fi

  if grep -qiE "$AUTH_ERROR_PATTERN" "$err_file"; then
    AUTH_ERROR=1
  fi
  rm -f "$err_file"
  return 0
}

# Main logic
if use_relay; then
  if token=$(get_token); then
    export GH_TOKEN="$token"
  fi
fi

# Find original gh binary
GH_BIN=""
if [[ -x /usr/bin/gh ]]; then
  GH_BIN=/usr/bin/gh
else
  echo "gh-relay: Error: Cannot find gh binary at /usr/bin/gh" >&2
  exit 1
fi

if run_gh "$@"; then
  exit 0
fi
status=$?

# If auth failed and we're in relay mode, refresh token and retry once.
if use_relay && [[ "$AUTH_ERROR" -eq 1 ]]; then
  clear_cached_token
  if token=$(fetch_fresh_token); then
    export GH_TOKEN="$token"
    if run_gh "$@"; then
      exit 0
    fi
    status=$?
  fi
fi

exit "$status"
