#!/usr/bin/env bash
# 08-memorybank-setup.sh — Create Vertex AI Agent Engine (Memory Bank) instance
# Uses REST API directly — no SDK version dependency
# Idempotent — safe to run multiple times
# Source .env before running: source .env

set -euo pipefail

: "${PROJECT_ID:?Set PROJECT_ID}"
: "${REGION:?Set REGION}"
: "${AGENT_SA:?Set AGENT_SA}"

MEMORY_INSTANCE_DISPLAY_NAME="${MEMORY_INSTANCE_DISPLAY_NAME:-${PROJECT_ID}-memory-bank}"

# ── Region Guard ────────────────────────────────────────────────────────
source "$(dirname "$0")/99-region-guard.sh"
MEMORY_REGION=$(validate_agent_engine_region "${MEMORY_REGION:-${REGION}}")
log_region_context "Agent-Engine" "$MEMORY_REGION" "VALIDATED"

echo "=== 08: Setting up Vertex AI Memory Bank (Agent Engine) ==="
echo "    Project  : ${PROJECT_ID}"
echo "    Region   : ${MEMORY_REGION}"
echo "    Instance : ${MEMORY_INSTANCE_DISPLAY_NAME}"
echo ""

bash "$(dirname "$0")/00-authcheck.sh"

gcloud services enable aiplatform.googleapis.com \
  --project="${PROJECT_ID}" --quiet
echo "✓ APIs enabled"

# ── REST: token + endpoint ────────────────────────────────────────────────────
TOKEN=$(gcloud auth print-access-token)
API="https://${MEMORY_REGION}-aiplatform.googleapis.com/v1beta1"
PARENT="projects/${PROJECT_ID}/locations/${MEMORY_REGION}"

# ── Check if instance exists ───────────────────────────────────────────────────
echo "  Checking for existing Agent Engine instance..."
EXISTING_LIST=$(curl -sf \
  -H "Authorization: Bearer ${TOKEN}" \
  "${API}/${PARENT}/reasoningEngines" 2>/dev/null || echo '{}')

ENGINE_NAME=$(echo "${EXISTING_LIST}" | python3 -c "
import sys, json
data = json.load(sys.stdin)
for e in data.get('reasoningEngines', []):
    if e.get('displayName') == '${MEMORY_INSTANCE_DISPLAY_NAME}':
        print(e['name'])
        break
" 2>/dev/null || true)

if [[ -n "${ENGINE_NAME}" ]]; then
  echo "✓ Memory Bank already exists: ${ENGINE_NAME}"
else
  echo "  Creating Agent Engine '${MEMORY_INSTANCE_DISPLAY_NAME}' in ${MEMORY_REGION}..."

  RESPONSE=$(curl -sf -X POST \
    -H "Authorization: Bearer ${TOKEN}" \
    -H "Content-Type: application/json" \
    "${API}/${PARENT}/reasoningEngines" \
    -d '{
      "displayName": "'"${MEMORY_INSTANCE_DISPLAY_NAME}"'",
      "spec": {}
    }' 2>/dev/null || echo '{}')

  OPERATION=$(echo "${RESPONSE}" | python3 -c "
import sys, json
d = json.load(sys.stdin)
print(d.get('name', ''))
" 2>/dev/null || true)

  if [[ -z "${OPERATION}" ]]; then
    echo "ERROR: No operation returned. Response:"
    echo "${RESPONSE}"
    exit 1
  fi

  echo "  Waiting for LRO: ${OPERATION}"
  for i in $(seq 1 30); do
    sleep 10
    LRO=$(curl -sf \
      -H "Authorization: Bearer ${TOKEN}" \
      "https://${MEMORY_REGION}-aiplatform.googleapis.com/v1beta1/${OPERATION}" 2>/dev/null || echo '{}')
    DONE=$(echo "${LRO}" | python3 -c "
import sys, json
print(json.load(sys.stdin).get('done', False))
" 2>/dev/null || echo 'False')
    if [[ "${DONE}" == "True" ]]; then
      # Re-list to get name reliably
      ENGINE_NAME=$(curl -sf \
        -H "Authorization: Bearer ${TOKEN}" \
        "${API}/${PARENT}/reasoningEngines" 2>/dev/null \
        | python3 -c "
import sys, json
data = json.load(sys.stdin)
for e in data.get('reasoningEngines', []):
    if e.get('displayName') == '${MEMORY_INSTANCE_DISPLAY_NAME}':
        print(e['name'])
        break
" 2>/dev/null || true)
      echo "✓ Agent Engine created: ${ENGINE_NAME}"
      break
    fi
    echo "  ... still waiting (${i}/30)"
  done

  if [[ -z "${ENGINE_NAME}" ]]; then
    echo "ERROR: Agent Engine creation timed out or failed."
    exit 1
  fi
fi

# ── Write output ──────────────────────────────────────────────────────────────────
echo "${ENGINE_NAME}" > /tmp/memory_engine_name.txt

echo ""
echo "  Resource name : ${ENGINE_NAME}"
echo "  Region        : ${MEMORY_REGION}"
echo ""
echo "  ACTION REQUIRED — add to .env:"
echo "    export MEMORY_ENGINE_NAME=\"${ENGINE_NAME}\""
echo ""
echo "=== 08: Memory Bank setup COMPLETE ==="
echo "    NEXT: copy MEMORY_ENGINE_NAME into .env, then run 05-cloudrun-deploy.sh"
echo ""