#!/bin/sh
# context7 - query library/framework documentation via the Context7 public
#           MCP server (https://mcp.context7.com/mcp).
#
# Context7 keeps up-to-date docs and code examples for thousands of libraries
# and frameworks and exposes them through an MCP (Model Context Protocol)
# server. This script speaks the JSON-RPC MCP protocol over HTTP:
#   1. initialize                  -> obtains an Mcp-Session-Id (optional)
#   2. notifications/initialized   -> no reply
#   3. tools/call                  -> runs a Context7 tool and prints text
#
# Context7 provides two tools:
#   * resolve-library-id  maps a library name to a Context7 ID (/org/project).
#   * query-docs          fetches docs + examples for a resolved library ID.
#
# Usage:
#   context7 resolve <libraryName> [query]
#       Search for a library and print the candidate Context7 IDs with their
#       description, snippet count, reputation and benchmark score.
#
#   context7 query <libraryId> <query>
#       Fetch documentation for an explicit library ID (format /org/project
#       or /org/project/version). The ID usually comes from `resolve`.
#
#   context7 docs <libraryName> <query>
#       Convenience: resolve the library, auto-pick the top-ranked match and
#       immediately query its documentation.
#
#   context7 --help
#
# Environment:
#   CONTEXT7_MCP_ENDPOINT  optional; default https://mcp.context7.com/mcp
#   CONTEXT7_API_KEY       optional; if set, sent as the X-Context7-API-Key
#                          request header (for higher rate limits / private
#                          docs). No extra headers are added when unset.
#
# Dependencies: curl, jq.

set -eu

usage() {
  sed -n '/^# Usage:/,/^# Dependencies:/p' "$0" | sed 's/^# \{0,1\}//'
}

if [ "$#" -lt 1 ]; then
  usage >&2
  exit 1
fi

CMD="$1"
shift

case "$CMD" in
  -h|--help|help)
    usage
    exit 0
    ;;
  resolve|query|docs)
    ;;
  *)
    echo "context7: unknown command '$CMD' (try --help)" >&2
    exit 1
    ;;
esac

# --- Validate argument counts per command. ---
case "$CMD" in
  resolve)
    if [ "$#" -lt 1 ]; then
      echo "context7: resolve requires <libraryName> [query]" >&2
      exit 1
    fi
    LIB_NAME="$1"
    QUERY="${2:-$1}"
    ;;
  query)
    if [ "$#" -lt 2 ]; then
      echo "context7: query requires <libraryId> <query>" >&2
      exit 1
    fi
    LIB_ID="$1"
    QUERY="$2"
    ;;
  docs)
    if [ "$#" -lt 2 ]; then
      echo "context7: docs requires <libraryName> <query>" >&2
      exit 1
    fi
    LIB_NAME="$1"
    QUERY="$2"
    ;;
esac

ENDPOINT="${CONTEXT7_MCP_ENDPOINT:-https://mcp.context7.com/mcp}"

# Optional auth header (single token, no spaces expected in an API key).
AUTH_ARGS=""
if [ -n "${CONTEXT7_API_KEY:-}" ]; then
  AUTH_ARGS="-H X-Context7-API-Key:${CONTEXT7_API_KEY}"
fi

# --- MCP handshake + tool call helper. ---
# Args: $1 = tool name, $2 = arguments JSON. Prints the parsed text content.
mcp_call() {
  TOOL="$1"
  ARGS="$2"

  INIT_HEADERS=$(mktemp)
  INIT_BODY=$(mktemp)
  curl -s -D "$INIT_HEADERS" -X POST "$ENDPOINT" \
    -H "Content-Type: application/json" \
    -H "Accept: application/json, text/event-stream" \
    $AUTH_ARGS \
    -d '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"context7","version":"1.0"}}}' \
    > "$INIT_BODY"

  SID=$(grep -i '^mcp-session-id:' "$INIT_HEADERS" | tr -d '\r' | awk '{print $2}')

  if [ -n "$SID" ]; then
    curl -s -X POST "$ENDPOINT" \
      -H "Content-Type: application/json" \
      -H "Accept: application/json, text/event-stream" \
      -H "Mcp-Session-Id: $SID" \
      $AUTH_ARGS \
      -d '{"jsonrpc":"2.0","method":"notifications/initialized"}' > /dev/null
    SID_ARGS="-H Mcp-Session-Id:$SID"
  else
    SID_ARGS=""
  fi

  CALL_BODY=$(mktemp)
  curl -s -X POST "$ENDPOINT" \
    -H "Content-Type: application/json" \
    -H "Accept: application/json, text/event-stream" \
    $SID_ARGS \
    $AUTH_ARGS \
    -d "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/call\",\"params\":{\"name\":\"$TOOL\",\"arguments\":$ARGS}}" \
    > "$CALL_BODY"

  # Take the last "data:" SSE payload (or the whole body if plain JSON).
  DATA=$(grep -E '^data:[[:space:]]' "$CALL_BODY" | tail -n1 | sed 's/^data:[[:space:]]*//')
  if [ -z "$DATA" ]; then
    DATA=$(cat "$CALL_BODY")
  fi

  rm -f "$INIT_HEADERS" "$INIT_BODY" "$CALL_BODY"

  if [ -z "$DATA" ]; then
    echo "context7: no response from MCP server" >&2
    return 1
  fi

  # Surface protocol errors.
  ERROR_MSG=$(printf '%s' "$DATA" | jq -r 'if .error then (.error|tostring) else empty end')
  if [ -n "$ERROR_MSG" ]; then
    echo "context7 error: $ERROR_MSG" >&2
    return 1
  fi

  printf '%s' "$DATA" | jq -r '
    (.result.content // [])
    | map(select(.type == "text"))
    | map(.text)
    | .[]
  '
}

# --- Dispatch. ---
case "$CMD" in
  resolve)
    ARGS=$(jq -n --arg n "$LIB_NAME" --arg q "$QUERY" \
      '{libraryName:$n, query:$q}')
    mcp_call "resolve-library-id" "$ARGS"
    ;;
  query)
    ARGS=$(jq -n --arg id "$LIB_ID" --arg q "$QUERY" \
      '{libraryId:$id, query:$q}')
    mcp_call "query-docs" "$ARGS"
    ;;
  docs)
    RARGS=$(jq -n --arg n "$LIB_NAME" --arg q "$QUERY" \
      '{libraryName:$n, query:$q}')
    RESOLVE_OUT=$(mcp_call "resolve-library-id" "$RARGS") || exit 1
    # Auto-pick the first candidate library ID from the resolve output.
    TOP_ID=$(printf '%s\n' "$RESOLVE_OUT" \
      | grep -m1 -oE '/[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+' || true)
    if [ -z "$TOP_ID" ]; then
      echo "context7: could not extract a library ID from resolve output:" >&2
      printf '%s\n' "$RESOLVE_OUT" >&2
      exit 1
    fi
    echo "context7: using library ID $TOP_ID (from resolve of '$LIB_NAME')" >&2
    QARGS=$(jq -n --arg id "$TOP_ID" --arg q "$QUERY" \
      '{libraryId:$id, query:$q}')
    mcp_call "query-docs" "$QARGS"
    ;;
esac

exit 0