#!/usr/bin/env bash
# RocketLit Claude setup. Re-runnable: it replaces its own block rather than
# appending a second copy.
set -euo pipefail

GATEWAY='https://vertex.litrocket.dev'
BEGIN='# >>> rocketlit claude gateway >>>'
END='# <<< rocketlit claude gateway <<<'

die() { echo "$*" >&2; exit 1; }

KEY="${1:-}"
if [ -z "$KEY" ]; then
  printf 'Paste your gateway key (sk-...): ' >&2
  read -rs KEY
  printf '\n' >&2
fi

case "$KEY" in
  sk-*) ;;
  *) die "That does not look like a gateway key. It starts with sk-." ;;
esac

case "$(uname -s)" in
  Darwin) os=mac ;;
  Linux) os=linux ;;
  *) die "Unsupported system: $(uname -s). Ask in #engineering." ;;
esac

# Where the exports go depends on the login shell, not the one running this
# script, since that is the shell that will read them next time.
case "$(basename "${SHELL:-}")" in
  zsh) rc="$HOME/.zshrc" ;;
  bash) if [ "$os" = mac ]; then rc="$HOME/.bash_profile"; else rc="$HOME/.bashrc"; fi ;;
  *) die "Only zsh and bash are handled. Add the exports by hand from the portal." ;;
esac

echo "Checking your key"
info="$(curl -fsSL -H "Authorization: Bearer $KEY" "$GATEWAY/key/info" 2>/dev/null)" \
  || die "The gateway rejected that key. Mint a fresh one from the portal."

# One JSON field, no jq: split on commas so each line holds a single key.
email="$(printf '%s' "$info" | tr ',' '\n' | sed -n 's/.*"user_id"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p' | head -1)"

# Pins come from what the key can actually call, not from a tier written down
# here. A tier without opus degrades to sonnet on purpose: leaving the default
# unset would let Claude Code fall back to a model the key cannot call.
models="$(curl -fsSL -H "Authorization: Bearer $KEY" "$GATEWAY/v1/models" 2>/dev/null || true)"
case "$models" in
  *'"claude-opus-5"'*) primary='claude-opus-5[1m]' ;;
  *) primary='claude-sonnet-5[1m]' ;;
esac

echo "Writing gateway settings to $rc"

tmp="$(mktemp)"
if [ -f "$rc" ]; then
  # Drop any previous block, keep everything else, then append the current one.
  awk -v b="$BEGIN" -v e="$END" 'index($0,b){skip=1} !skip{print} index($0,e){skip=0}' "$rc" > "$tmp"
fi

{
  echo "$BEGIN"
  echo "export ANTHROPIC_BASE_URL=$GATEWAY"
  echo "export ANTHROPIC_AUTH_TOKEN=$KEY"
  echo "export ANTHROPIC_MODEL=$primary"
  echo "export ANTHROPIC_DEFAULT_OPUS_MODEL=$primary"
  echo "export ANTHROPIC_DEFAULT_SONNET_MODEL=claude-sonnet-5[1m]"
  echo "export ANTHROPIC_DEFAULT_HAIKU_MODEL=claude-haiku-4-5"
  echo "export CLAUDE_CODE_ENABLE_TELEMETRY=1"
  echo "export OTEL_METRICS_EXPORTER=otlp"
  echo "export OTEL_LOGS_EXPORTER=otlp"
  echo "export OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf"
  echo "export OTEL_EXPORTER_OTLP_ENDPOINT=https://setup.litrocket.dev"
  echo "export OTEL_EXPORTER_OTLP_HEADERS=\"Authorization=Bearer $KEY\""
  [ -n "$email" ] && echo "export OTEL_RESOURCE_ATTRIBUTES=user.email=$email"
  echo "$END"
} >> "$tmp"

mv "$tmp" "$rc"

if command -v claude >/dev/null 2>&1; then
  echo "Adding the RocketLit MCP server"
  # Remove first: adding a name that already exists is an error, and this
  # script is meant to be safe to run twice.
  claude mcp remove --scope user rocketlit >/dev/null 2>&1 || true
  claude mcp add --scope user --transport http rocketlit https://vertex.litrocket.dev/mcp/ \
  --header "x-litellm-api-key: Bearer ${KEY}" \
  --header "x-mcp-servers: rocketlit"
else
  echo "Claude Code is not on PATH, skipping the MCP server. Install it, then re-run this."
fi

if [ "$os" = mac ]; then
  profile="$HOME/Downloads/rocketlit-claude.mobileconfig"
  cat > "$profile" <<EOF
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
  <key>PayloadContent</key>
  <array>
    <dict>
      <key>PayloadType</key>
      <string>com.anthropic.claudefordesktop</string>
      <key>PayloadIdentifier</key>
      <string>com.rocketlit.llm-proxy.gateway</string>
      <key>PayloadUUID</key>
      <string>1f1c3b8e-6a1a-4f7e-9b3d-5c2a7e4d8f10</string>
      <key>PayloadVersion</key>
      <integer>1</integer>
      <key>PayloadDisplayName</key>
      <string>RocketLit Claude Gateway</string>
      <key>inferenceProvider</key>
      <string>gateway</string>
      <key>inferenceGatewayBaseUrl</key>
      <string>https://vertex.litrocket.dev</string>
      <key>inferenceGatewayApiKey</key>
      <string>${KEY}</string>
      <key>inferenceGatewayAuthScheme</key>
      <string>bearer</string>
      <key>managedMcpServers</key>
      <array>
        <dict>
          <key>name</key>
          <string>rocketlit</string>
          <key>transport</key>
          <string>http</string>
          <key>url</key>
          <string>https://mcp.rocketlit.com/mcp</string>
          <key>oauth</key>
          <dict>
            <key>authorizationServer</key>
            <array>
              <string>https://mcp.rocketlit.com</string>
            </array>
          </dict>
        </dict>
      </array>
    </dict>
  </array>
  <key>PayloadDisplayName</key>
  <string>RocketLit Claude Gateway</string>
  <key>PayloadIdentifier</key>
  <string>com.rocketlit.llm-proxy</string>
  <key>PayloadType</key>
  <string>Configuration</string>
  <key>PayloadUUID</key>
  <string>2d4e6f80-9c1b-4a2e-8d5f-3b7c1e9a0d24</string>
  <key>PayloadVersion</key>
  <integer>1</integer>
</dict>
</plist>
EOF
  echo "Opening the Claude Desktop profile"
  # Installing a profile needs a human in System Settings; this only presents it.
  open "$profile"
  echo "Approve it under System Settings > General > Device Management, then quit Claude Desktop fully and reopen it."
fi

echo
echo "Done. Open a new terminal, or run: source $rc"