#!/bin/bash
#
# Nexus Memory - Optimized Recall Memory Hook (GraphRAG v2)
# Recalls relevant context from Nexus GraphRAG with entity extraction,
# knowledge graph traversal, and episodic memory support.
#
# PERFORMANCE OPTIMIZED:
# - Query caching with configurable TTL
# - Fast dependency checks
# - Connection keep-alive for faster requests
# - Optimized JSON parsing
#
# Usage:
#   echo '{"query": "..."}' | recall-memory.sh
#
# Environment Variables:
#   NEXUS_API_KEY        - API key for authentication (REQUIRED)
#   NEXUS_API_URL        - API endpoint (default: https://api.adverant.ai)
#   NEXUS_COMPANY_ID     - Company identifier (default: adverant)
#   NEXUS_APP_ID         - Application identifier (default: claude-code)
#   NEXUS_VERBOSE        - Set to 1 for debug output
#
# GraphRAG Enhancement Options:
#   NEXUS_INCLUDE_EPISODIC   - Include episodic memory context (default: true)
#   NEXUS_INCLUDE_DOCUMENTS  - Include document context (default: false for speed)
#   NEXUS_INCLUDE_ENTITIES   - Include entity information (default: true)
#   NEXUS_INCLUDE_FACTS      - Include extracted facts (default: true)
#   NEXUS_INCLUDE_FOLLOWUPS  - Include suggested follow-ups (default: false)
#   NEXUS_GRAPH_DEPTH        - Multi-hop graph traversal depth (default: 1)
#   NEXUS_MAX_TOKENS         - Token budget for recall (default: 3000)
#   NEXUS_CACHE_TTL          - Cache TTL in seconds (default: 600)
#
# Output:
#   Rich JSON object with memories, documents, entities, facts, and suggestions
#

set -o pipefail

# Source the API key helper for interactive prompting
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/api-key-helper.sh" 2>/dev/null || true

# Configuration with environment variable overrides
NEXUS_API_KEY="${NEXUS_API_KEY:-}"
NEXUS_API_URL="${NEXUS_API_URL:-https://api.adverant.ai}"
COMPANY_ID="${NEXUS_COMPANY_ID:-adverant}"
APP_ID="${NEXUS_APP_ID:-claude-code}"
VERBOSE="${NEXUS_VERBOSE:-0}"

# GraphRAG Enhancement Configuration (optimized defaults)
INCLUDE_EPISODIC="${NEXUS_INCLUDE_EPISODIC:-true}"
INCLUDE_DOCUMENTS="${NEXUS_INCLUDE_DOCUMENTS:-false}"  # Disabled for speed
INCLUDE_ENTITIES="${NEXUS_INCLUDE_ENTITIES:-true}"
INCLUDE_FACTS="${NEXUS_INCLUDE_FACTS:-true}"
INCLUDE_FOLLOWUPS="${NEXUS_INCLUDE_FOLLOWUPS:-false}"  # Disabled for speed
GRAPH_DEPTH="${NEXUS_GRAPH_DEPTH:-1}"  # Reduced for speed
MAX_TOKENS="${NEXUS_MAX_TOKENS:-3000}"

# Cache configuration
CACHE_DIR="${HOME}/.claude/session-env/recall-cache"
CACHE_TTL="${NEXUS_CACHE_TTL:-600}"  # 10 minutes
SKIP_CACHE="${NEXUS_SKIP_CACHE:-0}"

# Logging function
log() {
  if [[ "$VERBOSE" == "1" ]]; then
    echo "[recall-memory] $1" >&2
  fi
}

log_error() {
  echo "[recall-memory] ERROR: $1" >&2
}

# Default empty response with enhanced structure
EMPTY_RESPONSE='{
  "memories": [],
  "documents": [],
  "entities": [],
  "facts": [],
  "episodic_context": [],
  "suggested_followups": []
}'

# =========================================================
# FAST PATH: Early exit checks
# =========================================================

# Check for API key (REQUIRED) - Interactive prompt if running in terminal
# recall-memory.sh can be user-invoked or called by other hooks
if [[ -z "$NEXUS_API_KEY" ]]; then
  # Try loading from saved file first
  if type load_api_key &>/dev/null; then
    load_api_key
  fi

  # If still no key and we're interactive, prompt
  if [[ -z "$NEXUS_API_KEY" ]] && [[ -t 0 ]]; then
    if type require_api_key &>/dev/null; then
      require_api_key --interactive || {
        echo "$EMPTY_RESPONSE"
        exit 1
      }
    else
      log_error "NEXUS_API_KEY not set"
      echo "$EMPTY_RESPONSE"
      exit 1
    fi
  elif [[ -z "$NEXUS_API_KEY" ]]; then
    log_error "NEXUS_API_KEY not set"
    echo "$EMPTY_RESPONSE"
    exit 1
  fi
fi

# Fast dependency check
type jq &>/dev/null || { log_error "jq not available"; echo "$EMPTY_RESPONSE"; exit 1; }
type curl &>/dev/null || { log_error "curl not available"; echo "$EMPTY_RESPONSE"; exit 1; }

# Ensure cache directory exists
mkdir -p "$CACHE_DIR" 2>/dev/null

# Read input from stdin
INPUT=$(cat)

if [[ -z "$INPUT" ]]; then
  log "No input provided"
  echo "$EMPTY_RESPONSE"
  exit 0
fi

log "Received input: ${INPUT:0:100}..."

# =========================================================
# INPUT PARSING
# =========================================================

# Extract query and limit with single jq call
read -r QUERY LIMIT < <(echo "$INPUT" | jq -r '[
  (.query // .prompt // .tool_input.command // "recent context"),
  (.limit // 10)
] | @tsv' 2>/dev/null || echo "recent context	10")

# Validate limit
if ! [[ "$LIMIT" =~ ^[0-9]+$ ]]; then
  LIMIT=10
fi

# Get optional filters from input
FILTERS=$(echo "$INPUT" | jq -c '.filters // {}' 2>/dev/null)
if [[ "$FILTERS" == "null" ]] || [[ -z "$FILTERS" ]]; then
  FILTERS="{}"
fi

# Get current project for context
PROJECT_NAME=$(basename "$(pwd)")
PROJECT_DIR=$(pwd)

log "Query: $QUERY"
log "Limit: $LIMIT"
log "Project: $PROJECT_NAME"

# Skip if no query
if [[ -z "$QUERY" ]] || [[ "$QUERY" == "null" ]]; then
  log "No query provided"
  echo "$EMPTY_RESPONSE"
  exit 0
fi

# =========================================================
# CACHE CHECK
# =========================================================

# Generate query hash
QUERY_KEY="${PROJECT_NAME}:${QUERY:0:150}:$LIMIT"
QUERY_HASH=$(echo -n "$QUERY_KEY" | md5 2>/dev/null || echo -n "$QUERY_KEY" | md5sum 2>/dev/null | cut -d' ' -f1)
CACHE_FILE="${CACHE_DIR}/recall-${QUERY_HASH}"

# Check cache
if [[ "$SKIP_CACHE" != "1" ]] && [[ -f "$CACHE_FILE" ]]; then
  CACHE_AGE=$(( $(date +%s) - $(stat -f %m "$CACHE_FILE" 2>/dev/null || stat -c %Y "$CACHE_FILE" 2>/dev/null || echo 0) ))

  if [[ "$CACHE_AGE" -lt "$CACHE_TTL" ]]; then
    log "Cache hit (age: ${CACHE_AGE}s)"
    cat "$CACHE_FILE"
    exit 0
  else
    log "Cache expired"
    rm -f "$CACHE_FILE" 2>/dev/null
  fi
fi

# =========================================================
# API REQUEST
# =========================================================

# Build the optimized payload for GraphRAG retrieval
PAYLOAD=$(jq -n \
  --arg query "$QUERY" \
  --argjson limit "$LIMIT" \
  --arg project "$PROJECT_NAME" \
  --arg projectDir "$PROJECT_DIR" \
  --argjson filters "$FILTERS" \
  --argjson includeEpisodic "$INCLUDE_EPISODIC" \
  --argjson includeDocuments "$INCLUDE_DOCUMENTS" \
  --argjson includeEntities "$INCLUDE_ENTITIES" \
  --argjson includeFacts "$INCLUDE_FACTS" \
  --argjson includeFollowups "$INCLUDE_FOLLOWUPS" \
  --argjson graphDepth "$GRAPH_DEPTH" \
  --argjson maxTokens "$MAX_TOKENS" \
  '{
    query: $query,
    limit: $limit,
    context: {
      project: $project,
      projectDir: $projectDir
    },
    filters: ($filters + {project: $project}),
    include_episodic: $includeEpisodic,
    include_documents: $includeDocuments,
    include_entities: $includeEntities,
    include_facts: $includeFacts,
    include_followups: $includeFollowups,
    enable_graph_traversal: ($graphDepth > 0),
    graph_depth: $graphDepth,
    max_tokens: $maxTokens,
    hybrid_search: true,
    rerank: true,
    fast_mode: true
  }')

# Use the UNIFIED /api/memory endpoint for all memory operations
# This single endpoint handles both store (with content) and recall (with query)
# Returns: unified_memories, document_context, episodic_context, entities, facts
ENDPOINT="$NEXUS_API_URL/api/memory"
log "Recalling from unified endpoint: $ENDPOINT"

# Make request with optimized settings
RESPONSE=$(curl -s -w "\n%{http_code}" -X POST "$ENDPOINT" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $NEXUS_API_KEY" \
  -H "X-Company-ID: $COMPANY_ID" \
  -H "X-App-ID: $APP_ID" \
  -H "X-User-ID: ${USER:-unknown}" \
  -H "Connection: keep-alive" \
  -d "$PAYLOAD" \
  --connect-timeout 3 \
  --max-time 10 2>&1)

# Parse response
HTTP_CODE=$(echo "$RESPONSE" | tail -n1)
BODY=$(echo "$RESPONSE" | sed '$d')

log "Response code: $HTTP_CODE"

# Handle API errors
if [[ "$HTTP_CODE" != "200" ]]; then
  if [[ "$VERBOSE" == "1" ]]; then
    echo "❌ Recall failed (HTTP $HTTP_CODE)" >&2

    ERROR_MSG=$(echo "$BODY" | jq -r '.error.message // .message // ""' 2>/dev/null)
    if [[ -n "$ERROR_MSG" ]]; then
      echo "   Error: $ERROR_MSG" >&2
    fi

    if [[ "$HTTP_CODE" == "401" ]]; then
      echo "   Check: NEXUS_API_KEY environment variable" >&2
    fi
  fi

  # Return empty response so Claude doesn't get errors
  echo "$EMPTY_RESPONSE"
  exit 1  # Non-blocking error for logging
fi

# Validate JSON response
if ! echo "$BODY" | jq -e . &>/dev/null; then
  log_error "Invalid JSON response"
  echo "$EMPTY_RESPONSE"
  exit 0
fi

# =========================================================
# RESPONSE NORMALIZATION
# =========================================================

# Normalize response to expected structure - single jq pass
# Handle multiple response formats:
# 1. Gateway wrapper: { success: true, data: { unified_memories: [...], ... } }
# 2. Direct response: { unified_memories: [...], ... }
# 3. Enhanced format: { data: { results: { memories: [...] } } }
NORMALIZED=$(echo "$BODY" | jq '
  # Unwrap gateway/API layers to get to actual data
  (
    if .data.data then .data.data      # Double-wrapped gateway response
    elif .data.unified_memories then .data  # Single-wrapped with unified_memories
    elif .data.results then .data.results   # Old format with results wrapper
    elif .data then .data                   # Single-wrapped generic
    else .                                  # Direct response
    end
  ) as $result |
  {
    memories: ($result.unified_memories // $result.memories // []),
    documents: ($result.document_context // $result.documents // []),
    entities: ($result.entities_mentioned // $result.entities // []),
    facts: ($result.relevant_facts // $result.facts // []),
    episodic_context: ($result.episodic_context // []),
    suggested_followups: ($result.suggested_followups // []),
    query_understanding: ($result.query_understanding // null),
    confidence_score: ($result.confidence_score // null),
    context: {
      project: ($result.context.project // "unknown"),
      query: ($result.query // $result.metadata.query // "")
    }
  }
')

# Cache the result
if [[ "$SKIP_CACHE" != "1" ]]; then
  echo "$NORMALIZED" > "$CACHE_FILE" 2>/dev/null
  log "Cached result for ${CACHE_TTL}s"

  # Cleanup old cache files (async)
  (find "$CACHE_DIR" -type f -name "recall-*" -mmin +$((CACHE_TTL / 60 + 1)) -delete 2>/dev/null &)
fi

# Output verbose info if enabled
if [[ "$VERBOSE" == "1" ]]; then
  MEMORY_COUNT=$(echo "$NORMALIZED" | jq '.memories | length')

  if [[ "$MEMORY_COUNT" == "0" ]]; then
    echo "⚠️  No memories found for query" >&2
  else
    echo "✅ Retrieved $MEMORY_COUNT memories" >&2
  fi
fi

echo "$NORMALIZED"