#!/bin/sh
# websearch - simple web search for Bantam via Exa's MCP server (Streamable HTTP).
#
# Talks to https://mcp.exa.ai/mcp using the JSON-RPC MCP protocol:
#   1. initialize                 -> obtains an Mcp-Session-Id from response headers
#   2. notifications/initialized  -> no reply
#   3. tools/call (web_search_exa) -> prints the formatted result text
#
# Usage:
#   websearch "your query"
#   websearch "your query" [num_results]
#
# Environment:
#   EXA_MCP_ENDPOINT  optional; default https://mcp.exa.ai/mcp
#   EXA_API_KEY       optional; if set, sent as an URL query parameter (?api_key=)
#                    so that NO extra HTTP headers are added beyond Content-Type.
#
# Dependencies: curl, jq.

set -eu

if [ "$#" -lt 1 ]; then
  echo "Usage: websearch \"query\" [num_results]" >&2
  exit 1
fi

QUERY="$1"
NUM="${2:-5}"

ENDPOINT="${EXA_MCP_ENDPOINT:-https://mcp.exa.ai/mcp}"

# Append API key as a query parameter (no extra headers).
if [ -n "${EXA_API_KEY:-}" ]; then
  URL="${ENDPOINT}?api_key=${EXA_API_KEY}"
else
  URL="${ENDPOINT}"
fi

# Build the tools/call arguments JSON safely (proper query escaping).
ARGS=$(jq -n --arg q "$QUERY" --argjson n "$NUM" \
  '{query:$q, numResults:$n}')

# --- Step 1: initialize, capture headers (for Mcp-Session-Id) and body. ---
INIT_HEADERS=$(mktemp)
INIT_BODY=$(mktemp)
curl -s -D "$INIT_HEADERS" -X POST "$URL" \
  -H "Content-Type: application/json" \
  -H "Accept: application/json, text/event-stream" \
  -d '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"websearch","version":"1.0"}}}' \
  > "$INIT_BODY"

SID=$(grep -i '^mcp-session-id:' "$INIT_HEADERS" | tr -d '\r' | awk '{print $2}')
if [ -z "$SID" ]; then
  echo "websearch: failed to obtain MCP session id" >&2
  cat "$INIT_BODY" >&2
  rm -f "$INIT_HEADERS" "$INIT_BODY"
  exit 1
fi

# --- Step 2: send initialized notification (fire and forget). ---
curl -s -X POST "$URL" \
  -H "Content-Type: application/json" \
  -H "Accept: application/json, text/event-stream" \
  -H "Mcp-Session-Id: $SID" \
  -d '{"jsonrpc":"2.0","method":"notifications/initialized"}' > /dev/null

# --- Step 3: call web_search_exa and extract the text result. ---
CALL_BODY=$(mktemp)
curl -s -X POST "$URL" \
  -H "Content-Type: application/json" \
  -H "Accept: application/json, text/event-stream" \
  -H "Mcp-Session-Id: $SID" \
  -d "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/call\",\"params\":{\"name\":\"web_search_exa\",\"arguments\":$ARGS}}" \
  > "$CALL_BODY"

# Parse the response: take the last "data:" SSE payload (or the whole body
# if it is plain JSON) and print the text blocks from the result.
DATA=$(grep -E '^data:[[:space:]]' "$CALL_BODY" | tail -n1 | sed 's/^data:[[:space:]]*//')
if [ -z "$DATA" ]; then
  DATA=$(cat "$CALL_BODY")
fi

if [ -z "$DATA" ]; then
  echo "websearch: no response from MCP server" >&2
  rm -f "$INIT_HEADERS" "$INIT_BODY" "$CALL_BODY"
  exit 1
fi

# Surface server/protocol errors, then emit text content.
ERROR_MSG=$(printf '%s' "$DATA" | jq -r 'if .error then (.error|tostring) else empty end')
if [ -n "$ERROR_MSG" ]; then
  echo "websearch error: $ERROR_MSG" >&2
  rm -f "$INIT_HEADERS" "$INIT_BODY" "$CALL_BODY"
  exit 1
fi

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

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

exit 0