#!/usr/bin/env bash
# claude-vllm — run Claude Code against the homelab LLM gateway instead of api.anthropic.com.
#
# The gateway (LiteLLM in front of vLLM) already speaks the Anthropic Messages
# API on /v1/messages, so no bridge or translation layer is needed — Claude Code
# talks to it directly once ANTHROPIC_BASE_URL points there. This script exists
# only to find the endpoint, credential and model you have already configured
# for another agent, instead of making you paste four exports every time.
#
# Discovery order (first hit wins, per field):
#   1. environment already set (ANTHROPIC_BASE_URL / ANTHROPIC_AUTH_TOKEN / ANTHROPIC_MODEL)
#   2. flags (--provider, --model)
#   3. its own config: $XDG_CONFIG_HOME/mcpctl/claude-vllm.jsonc  (~/.config/...)
#   4. ~/.pi/agent        settings.json + models.json + auth.json
#   5. ~/.prime/agent     settings.json + models.json + auth.json
#
# The config file is shaped like opencode's — a `provider` map with
# `options.baseURL` / `options.apiKey` and a `models` map — but it is OUR file.
# Reading opencode's own config would mean a credential rotation there silently
# changing what Claude Code authenticates with, and would tie two tools' configs
# together for no reason. The pi/prime homes remain as a convenience fallback so
# `claude-vllm` works before you have written a config at all.
#
# NO KEY IS EVER STORED IN THIS SCRIPT. `apiKey` may be a literal (in a 0600
# config), or `\${ENV_VAR}` / a bare env var NAME, so the secret can live in your
# environment or a password manager instead of on disk.
#
# Usage:
#   claude-vllm                      # default provider + model, then exec claude
#   claude-vllm --model deepseek-v4-max
#   claude-vllm --provider itaz --model deepseek-v4-fast -- -p "summarise this repo"
#   claude-vllm --list               # show what it would use (never prints the key)
#   claude-vllm --print-env          # print the exports and exit (don't run claude)
#   claude-vllm --init --api-key ... # write the config file (0600); reads stdin if omitted
set -euo pipefail

PROVIDER=""
MODEL=""
LIST=0
INIT=0
PRINT_ENV=0
API_KEY_ARG=""
BASE_URL_ARG=""
CLAUDE_ARGS=()

while [ $# -gt 0 ]; do
  case "$1" in
    --provider) PROVIDER="${2:?--provider needs a value}"; shift 2 ;;
    --model)    MODEL="${2:?--model needs a value}"; shift 2 ;;
    --list)     LIST=1; shift ;;
    --init)     INIT=1; shift ;;
    --api-key)  API_KEY_ARG="${2:?--api-key needs a value}"; shift 2 ;;
    --base-url) BASE_URL_ARG="${2:?--base-url needs a value}"; shift 2 ;;
    --print-env) PRINT_ENV=1; shift ;;
    -h|--help)  sed -n '2,25p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
    --)         shift; CLAUDE_ARGS+=("$@"); break ;;
    *)          CLAUDE_ARGS+=("$1"); shift ;;
  esac
done

command -v jq >/dev/null || { echo "claude-vllm: jq is required" >&2; exit 1; }

PI_HOME="${PI_AGENT_HOME:-$HOME/.pi/agent}"
PRIME_HOME="${PRIME_AGENT_HOME:-$HOME/.prime/agent}"
CONFIG="${CLAUDE_VLLM_CONFIG:-${XDG_CONFIG_HOME:-$HOME/.config}/mcpctl/claude-vllm.jsonc}"

# jq tolerates the // comments opencode.jsonc may contain only after we strip
# them; harmless for strict JSON.
read_json() { [ -f "$1" ] && sed 's://[^"]*$::' "$1" | jq -c . 2>/dev/null || echo '{}'; }

# Resolve an apiKey field that may be a literal, \${ENV_VAR}, or a bare env var
# NAME (the form pi's models.json uses). Keeping the indirection means the
# secret can live in the environment rather than on disk.
resolve_key() {
  local raw="$1"
  [ -n "$raw" ] || return 0
  case "$raw" in
    '${'*'}') local n="${raw#\$\{}"; n="${n%\}}"; printf '%s' "${!n:-}" ;;
    # A bare ALL_CAPS token that names a set variable is an env var reference,
    # not a key: no real API key looks like that.
    [A-Z_][A-Z0-9_]*) if [ -n "${!raw:-}" ]; then printf '%s' "${!raw}"; else printf '%s' "$raw"; fi ;;
    *) printf '%s' "$raw" ;;
  esac
}

# ── discover ─────────────────────────────────────────────────────────────────
# Each agent home is tried in turn; the first one that yields a base URL wins,
# and the credential is taken from that same home so we never pair one gateway's
# URL with another's key.
discover_from_agent_home() {
  local home="$1" settings models auth provider base key model
  settings=$(read_json "$home/settings.json")
  models=$(read_json "$home/models.json")
  auth=$(read_json "$home/auth.json")

  provider="$PROVIDER"
  [ -n "$provider" ] || provider=$(jq -r '.defaultProvider // empty' <<<"$settings")
  [ -n "$provider" ] || return 1

  base=$(jq -r --arg p "$provider" '.providers[$p].baseUrl // empty' <<<"$models")
  [ -n "$base" ] || return 1
  # `apiKey` in models.json names an env var; the secret itself lives in auth.json.
  key=$(jq -r --arg p "$provider" '.[$p].key // empty' <<<"$auth")

  model="$MODEL"
  [ -n "$model" ] || model=$(jq -r '.defaultModel // empty' <<<"$settings")

  FOUND_PROVIDER="$provider"; FOUND_BASE="$base"; FOUND_KEY="$key"; FOUND_MODEL="$model"
  FOUND_MODELS=$(jq -r --arg p "$provider" '.providers[$p].models[]?.id' <<<"$models")
  # model -> context for ALL models, so --init records every limit rather than
  # only the active one (Claude Code assumes 200k for anything it lacks).
  FOUND_MODEL_LIMITS=$(jq -c --arg p "$provider" \
    '[.providers[$p].models[]? | select(.contextWindow) | {(.id): {limit:{context:.contextWindow}}}] | add // {}' <<<"$models")
  FOUND_CONTEXT=$(jq -r --arg p "$provider" --arg m "$model" \
    '.providers[$p].models[]? | select(.id==$m) | .contextWindow // empty' <<<"$models")
  FOUND_SOURCE="$home"
  return 0
}

# Our own config, shaped like opencode's but deliberately a separate file.
discover_from_config() {
  local cfg provider base key
  [ -f "$CONFIG" ] || return 1
  cfg=$(read_json "$CONFIG")
  provider="$PROVIDER"
  [ -n "$provider" ] || provider=$(jq -r '(.model // "") | split("/")[0] // empty' <<<"$cfg")
  [ -n "$provider" ] || provider=$(jq -r '.provider | keys[0] // empty' <<<"$cfg")
  [ -n "$provider" ] || return 1

  base=$(jq -r --arg p "$provider" '.provider[$p].options.baseURL // empty' <<<"$cfg")
  [ -n "$base" ] || return 1
  key=$(resolve_key "$(jq -r --arg p "$provider" '.provider[$p].options.apiKey // empty' <<<"$cfg")")

  FOUND_PROVIDER="$provider"; FOUND_BASE="$base"; FOUND_KEY="$key"
  FOUND_MODEL="${MODEL:-$(jq -r '(.model // "") | split("/")[1] // empty' <<<"$cfg")}"
  FOUND_MODELS=$(jq -r --arg p "$provider" '.provider[$p].models | keys[]?' <<<"$cfg")
  FOUND_MODEL_LIMITS=$(jq -c --arg p "$provider" '.provider[$p].models // {}' <<<"$cfg")
  FOUND_CONTEXT=$(jq -r --arg p "$provider" --arg m "$FOUND_MODEL" \
    '.provider[$p].models[$m].limit.context // empty' <<<"$cfg")
  FOUND_SOURCE="$CONFIG"
  return 0
}

FOUND_PROVIDER=""; FOUND_BASE=""; FOUND_KEY=""; FOUND_MODEL=""; FOUND_MODELS=""; FOUND_MODEL_LIMITS=""; FOUND_CONTEXT=""; FOUND_SOURCE=""
discover_from_config \
  || discover_from_agent_home "$PI_HOME" \
  || discover_from_agent_home "$PRIME_HOME" \
  || true

# ── --init: write the config file ────────────────────────────────────────────
if [ "$INIT" = 1 ]; then
  init_base="${BASE_URL_ARG:-${FOUND_BASE:-}}"
  if [ -z "$init_base" ]; then
    echo "claude-vllm --init: no endpoint known. Pass --base-url https://your-gateway/v1" >&2
    exit 1
  fi
  init_key="$API_KEY_ARG"
  if [ -z "$init_key" ]; then
    # Read from stdin so the key never lands in shell history or the process
    # table (where --api-key is visible to every user via `ps`).
    if [ -t 0 ]; then
      printf 'API key (input hidden, or pass ${ENV_VAR} to keep it out of the file): ' >&2
      read -rs init_key; echo >&2
    else
      read -r init_key || true
    fi
  fi
  [ -n "$init_key" ] || { echo "claude-vllm --init: no API key given" >&2; exit 1; }

  init_provider="${PROVIDER:-${FOUND_PROVIDER:-default}}"
  init_model="${MODEL:-${FOUND_MODEL:-}}"
  # Every model's limit, not just the active one — switching with --model must
  # not silently drop back to Claude Code's assumed 200k.
  init_models_json=$(
    if [ -n "$FOUND_MODELS" ]; then
      known="${FOUND_MODEL_LIMITS:-{\}}"
      for m in $FOUND_MODELS; do jq -n --arg m "$m" '{($m): {}}'; done \
        | jq -s --argjson known "$known" 'add // {} | . * $known'
    else
      jq -n --arg m "$init_model" 'if $m == "" then {} else {($m): {}} end'
    fi
  )

  mkdir -p "$(dirname "$CONFIG")"
  umask 077
  jq -n --arg p "$init_provider" --arg model "$init_model" --arg base "$init_base" \
        --arg key "$init_key" --argjson models "$init_models_json" '
    {
      "//": "mcpctl claude-vllm config. Shaped like opencode.jsonc, but its own file. apiKey may be a literal, \"${ENV_VAR}\", or a bare env var NAME.",
      model: (if $model == "" then null else "\($p)/\($model)" end),
      provider: { ($p): { options: { baseURL: $base, apiKey: $key }, models: $models } }
    } | del(..|nulls)' > "$CONFIG.tmp.$$"
  mv "$CONFIG.tmp.$$" "$CONFIG"
  chmod 600 "$CONFIG"
  echo "Wrote $CONFIG (0600)" >&2
  echo "  provider: $init_provider   endpoint: $init_base   model: ${init_model:-<gateway default>}" >&2
  case "$init_key" in
    '${'*'}'|[A-Z_][A-Z0-9_]*) echo "  apiKey: kept as an environment reference, not a literal" >&2 ;;
    *) echo "  apiKey: stored in the file — readable only by you" >&2 ;;
  esac
  exit 0
fi

BASE="${ANTHROPIC_BASE_URL:-$FOUND_BASE}"
KEY="${ANTHROPIC_AUTH_TOKEN:-${ANTHROPIC_API_KEY:-$FOUND_KEY}}"
MODEL_ID="${MODEL:-${ANTHROPIC_MODEL:-$FOUND_MODEL}}"

if [ "$LIST" = 1 ]; then
  echo "provider: ${FOUND_PROVIDER:-<none found>}   (from ${FOUND_SOURCE:-nowhere})"
  echo "config:   $CONFIG $([ -f "$CONFIG" ] && echo '(present)' || echo '(absent — using fallback; run --init)')"
  echo "endpoint: ${BASE:-<none found>}"
  echo "credential: $([ -n "$KEY" ] && echo "found (${KEY:0:10}…)" || echo "<none found>")"
  echo "default model: ${MODEL_ID:-<none>}${FOUND_CONTEXT:+  (context ${FOUND_CONTEXT})}"
  echo "models:"
  [ -n "$FOUND_MODELS" ] && printf '  %s\n' $FOUND_MODELS || echo "  <none listed>"
  exit 0
fi

if [ -z "$BASE" ]; then
  echo "claude-vllm: no LLM endpoint found." >&2
  echo "  Looked in $PI_HOME, $PRIME_HOME and $OC_CONFIG." >&2
  echo "  Set ANTHROPIC_BASE_URL, or configure a provider in one of those." >&2
  exit 1
fi
if [ -z "$KEY" ]; then
  echo "claude-vllm: found $BASE but no credential for '${FOUND_PROVIDER}'." >&2
  echo "  Set ANTHROPIC_AUTH_TOKEN, or add the key to ${FOUND_SOURCE}/auth.json." >&2
  exit 1
fi

# Claude Code appends /v1/messages itself, so the stored provider baseUrl's
# trailing /v1 (an OpenAI-style base) has to come off or requests go to
# /v1/v1/messages.
BASE="${BASE%/}"; BASE="${BASE%/v1}"

export ANTHROPIC_BASE_URL="$BASE"
# ANTHROPIC_AUTH_TOKEN only. Setting ANTHROPIC_API_KEY as well makes Claude Code
# warn that "auth may not work as expected" — it wants exactly one. AUTH_TOKEN
# is the right one for a third-party gateway (sent as `Authorization: Bearer`),
# and an inherited ANTHROPIC_API_KEY would otherwise take precedence over it, so
# clear it for the child process.
export ANTHROPIC_AUTH_TOKEN="$KEY"
unset ANTHROPIC_API_KEY
[ -n "$MODEL_ID" ] && export ANTHROPIC_MODEL="$MODEL_ID"
# Without a substitute, the background/summarisation calls ask the gateway for a
# real Haiku it does not serve, and every one of them 404s.
export ANTHROPIC_SMALL_FAST_MODEL="${ANTHROPIC_SMALL_FAST_MODEL:-${CLAUDE_VLLM_FAST_MODEL:-$MODEL_ID}}"
export ANTHROPIC_DEFAULT_HAIKU_MODEL="${ANTHROPIC_DEFAULT_HAIKU_MODEL:-$ANTHROPIC_SMALL_FAST_MODEL}"
# Claude Code only knows the context window of models it ships a table for, and
# assumes 200k for anything else — so a 393k model would auto-compact at half
# its capacity. Tell it the real number when the provider config states one.
if [ -n "${FOUND_CONTEXT:-}" ] && [ -z "${CLAUDE_CODE_MAX_CONTEXT_TOKENS:-}" ]; then
  export CLAUDE_CODE_MAX_CONTEXT_TOKENS="$FOUND_CONTEXT"
fi
# Beta headers the gateway does not implement make it reject otherwise fine
# requests.
export CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS="${CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS:-1}"
# Stops Claude Code phoning home about a non-Anthropic endpoint.
export CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC="${CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC:-1}"

if [ "$PRINT_ENV" = 1 ]; then
  for v in ANTHROPIC_BASE_URL ANTHROPIC_MODEL ANTHROPIC_SMALL_FAST_MODEL \
           ANTHROPIC_DEFAULT_HAIKU_MODEL CLAUDE_CODE_MAX_CONTEXT_TOKENS \
           CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS \
           CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC; do
    [ -n "${!v:-}" ] || continue
    printf 'export %s=%q\n' "$v" "${!v}"
  done
  # Never printed: the credential. Use --list to confirm one was found.
  echo 'export ANTHROPIC_AUTH_TOKEN=<redacted>'
  echo 'unset ANTHROPIC_API_KEY'
  exit 0
fi

command -v claude >/dev/null || { echo "claude-vllm: claude is not on PATH" >&2; exit 1; }
echo "claude-vllm: ${FOUND_PROVIDER:-custom} · ${ANTHROPIC_MODEL:-<gateway default>} · $ANTHROPIC_BASE_URL" >&2
exec claude "${CLAUDE_ARGS[@]}"