#!/usr/bin/env bash
# 07-rag-setup.sh — Create Vertex AI RAG Engine corpus and upload initial documents
# Idempotent
# Source: Vertex AI RAG Engine SDK (google-cloud-aiplatform >= 1.87)
# NOTE: RAG Engine in us-central1 requires allowlist access.
#       Contact: vertex-ai-rag-engine-support@google.com
#       Alternative: set REGION=us-east1 or us-east4 for immediate access.
# Source .env before running: source .env
set -euo pipefail

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

echo "=== 07: Setting up Vertex AI RAG Engine ==="

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

# Create GCS bucket for corpus source documents (idempotent)
CORPUS_BUCKET="${PROJECT_ID}-agent-corpus"
if ! gsutil ls -b "gs://${CORPUS_BUCKET}" &>/dev/null; then
  gsutil mb -l "${REGION}" -b on "gs://${CORPUS_BUCKET}"
  echo "✓ GCS corpus bucket created: gs://${CORPUS_BUCKET}"
else
  echo "✓ GCS corpus bucket exists: gs://${CORPUS_BUCKET}"
fi

# Upload seed documents if present
if [[ -d "docs/corpus-seed" ]]; then
  gsutil -m cp docs/corpus-seed/*.md "gs://${CORPUS_BUCKET}/seed/" 2>/dev/null || true
  echo "✓ Seed documents uploaded to gs://${CORPUS_BUCKET}/seed/"
fi

# Run Python to create/update corpus
python3 - << PYEOF
import os, sys
try:
    import vertexai
    from vertexai.preview import rag
except ImportError:
    print("ERROR: google-cloud-aiplatform not installed. Run: pip install google-cloud-aiplatform>=1.87.0")
    sys.exit(1)

PROJECT_ID = os.environ["PROJECT_ID"]
REGION = os.environ["REGION"]
CORPUS_DISPLAY_NAME = os.environ["RAG_CORPUS_DISPLAY_NAME"]
CORPUS_BUCKET = f"{PROJECT_ID}-agent-corpus"

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

corpus = None
for c in rag.list_corpora():
    if c.display_name == CORPUS_DISPLAY_NAME:
        corpus = c
        print(f"✓ RAG corpus already exists: {c.name}")
        break

if corpus is None:
    embedding_config = rag.EmbeddingModelConfig(
        publisher_model="publishers/google/models/text-embedding-005"
    )
    corpus = rag.create_corpus(
        display_name=CORPUS_DISPLAY_NAME,
        embedding_model_config=embedding_config,
    )
    print(f"✓ RAG corpus created: {corpus.name}")

gcs_uri = f"gs://{CORPUS_BUCKET}/seed/"
try:
    rag.import_files(
        corpus_name=corpus.name,
        paths=[gcs_uri],
        chunk_size=512,
        chunk_overlap=50,
        max_embedding_requests_per_min=900,
    )
    print(f"✓ Documents imported from {gcs_uri}")
except Exception as e:
    print(f"  WARNING: Document import skipped or failed: {e}")
    print("  Import manually via: https://console.cloud.google.com/vertex-ai/rag")

with open("/tmp/rag_corpus_name.txt", "w") as f:
    f.write(corpus.name)
print(f"\n  Corpus resource name: {corpus.name}")
print(f"  Add to .env: RAG_CORPUS_NAME={corpus.name}")
PYEOF

echo ""
echo "=== 07: RAG Engine setup COMPLETE ==="
echo "  View corpus: https://console.cloud.google.com/vertex-ai/rag?project=${PROJECT_ID}"