#!/usr/bin/env bash
# 08-memorybank-setup.sh — Create Vertex AI Agent Platform Memory Bank instance
# Idempotent
# Requires: google-cloud-aiplatform >= 1.111.0
# UNKNOWN: gcloud CLI equivalent for client.agent_engines.create() — SDK-only (OQ-01)
# 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 ==="

bash "$(dirname "$0")/00-authcheck.sh"
gcloud services enable aiplatform.googleapis.com --quiet

# Run Python to create/update Agent Engine instance with Memory Bank
python3 - << PYEOF
import os, sys
try:
    import vertexai
    from vertexai.preview import agent_engines
except ImportError:
    print("ERROR: google-cloud-aiplatform >= 1.111.0 required.")
    print("  Run: pip install 'google-cloud-aiplatform>=1.111.0'")
    sys.exit(1)

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

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

existing = list(client.agent_engines.list())
instance = None
for eng in existing:
    if hasattr(eng, "display_name") and eng.display_name == DISPLAY_NAME:
        instance = eng
        print(f"✓ Memory Bank instance already exists: {eng.api_resource.name}")
        break

if instance is None:
    memory_bank_config = {
        "memories_ttl_days": 30,
        "generate_memory_config": {
            "trigger_config": {"trigger_type": "ON_SESSION_END"}
        },
        "similarity_search_config": {"top_k": 5}
    }
    instance = client.agent_engines.create(
        display_name=DISPLAY_NAME,
        spec={"context_spec": {"memory_bank_config": memory_bank_config}}
    )
    print(f"✓ Memory Bank instance created: {instance.api_resource.name}")

instance_name = instance.api_resource.name
with open("/tmp/memory_bank_instance.txt", "w") as f:
    f.write(instance_name)
print(f"\n  Instance resource name: {instance_name}")
print(f"  Add to .env: MEMORY_BANK_INSTANCE={instance_name}")
PYEOF

# Grant agent SA memory permissions
gcloud projects add-iam-policy-binding "${PROJECT_ID}" \
  --member="serviceAccount:${AGENT_SA}" \
  --role="roles/aiplatform.user" \
  --quiet 2>/dev/null || true
echo "✓ roles/aiplatform.user granted to ${AGENT_SA}"

echo ""
echo "=== 08: Memory Bank setup COMPLETE ==="
echo "  View instances: https://console.cloud.google.com/vertex-ai/agents?project=${PROJECT_ID}"