#!/usr/bin/env bash
# 08-memorybank-setup.sh — Create Vertex AI Agent Engine (Memory Bank) instance
# 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}"

echo "=== 08: Setting up Vertex AI Memory Bank (Agent Engine) ==="
echo "    Project  : ${PROJECT_ID}"
echo "    Region   : ${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"

# ── Ensure correct SDK version is installed ────────────────────────────────
echo "  Ensuring google-cloud-aiplatform>=1.111.0 ..."
pip install --quiet --user "google-cloud-aiplatform>=1.111.0" 2>&1 | tail -2

# Use the user-local python so .local/lib packages are picked up
PYTHON="$(which python3)"
if python3 -c "import sys; sys.exit(0 if any('.local' in p for p in sys.path) else 1)" 2>/dev/null; then
  PYTHON="python3"
else
  # Force user site-packages
  PYTHON="python3 -c 'import sys, site; sys.path.insert(0, site.getusersitepackages()); exec(open("/dev/stdin").read())'"
fi

# Simple check: can we import agent_engines?
if ! python3 -c "import sys, site; sys.path.insert(0, site.getusersitepackages()); from vertexai.preview import agent_engines" 2>/dev/null; then
  echo "ERROR: agent_engines still not importable after install. Try:"
  echo "  pip install --upgrade google-cloud-aiplatform"
  exit 1
fi
echo "✓ Python SDK ready (agent_engines importable)"

# ── Create or find Memory Bank instance ───────────────────────────────────
python3 - << 'PYEOF'
import sys, site
sys.path.insert(0, site.getusersitepackages())

import os

try:
    import vertexai
    from vertexai.preview import agent_engines
except ImportError as e:
    print(f"ERROR: Failed to import agent_engines: {e}")
    sys.exit(1)

PROJECT_ID = os.environ["PROJECT_ID"]
REGION     = os.environ["REGION"]
DISP_NAME  = os.environ.get("MEMORY_INSTANCE_DISPLAY_NAME", f"{PROJECT_ID}-memory-bank")

vertexai.init(project=PROJECT_ID, location=REGION)

existing = None
try:
    for eng in agent_engines.list():
        if hasattr(eng, "display_name") and eng.display_name == DISP_NAME:
            existing = eng
            print(f"✓ Memory Bank already exists: {eng.resource_name}")
            break
except Exception as e:
    print(f"  WARNING: Could not list agent engines: {e}")

if existing is None:
    try:
        instance = agent_engines.create(display_name=DISP_NAME)
        print(f"✓ Memory Bank created: {instance.resource_name}")
        existing = instance
    except Exception as e:
        print(f"ERROR: Failed to create Memory Bank: {e}")
        sys.exit(1)

if existing:
    resource_name = existing.resource_name
    with open("/tmp/memory_engine_name.txt", "w") as f:
        f.write(resource_name)
    print(f"")
    print(f"  Resource name : {resource_name}")
    print(f"")
    print(f"  ACTION REQUIRED — add to .env:")
    print(f'    export MEMORY_ENGINE_NAME="{resource_name}"')
PYEOF

echo ""
echo "=== 08: Memory Bank setup COMPLETE ==="
echo "    NEXT: copy MEMORY_ENGINE_NAME into .env, then run 05-cloudrun-deploy.sh"
echo ""