#!/usr/bin/env bash
# 07-rag-setup.sh — Create Vertex AI RAG Engine corpus and upload initial documents
# Idempotent — safe to run multiple times
# Source: Vertex AI RAG Engine SDK (google-cloud-aiplatform >= 1.87)
#
# REGIONAL AVAILABILITY (OQ-02):
#   RAG Engine in us-central1 requires allowlist access.
#   This script auto-falls-back to us-east1 if us-central1 is not accessible.
#   To bypass: export RAG_REGION=us-east1 before running.
#
# 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}"

# RAG region fallback: prefer env override, then REGION, fallback to us-east1
RAG_REGION="${RAG_REGION:-${REGION}}"
RAG_FALLBACK_REGION="us-east1"

echo "=== 07: Setting up Vertex AI RAG Engine ==="
echo "    Project  : ${PROJECT_ID}"
echo "    Region   : ${RAG_REGION} (fallback: ${RAG_FALLBACK_REGION})"
echo "    Corpus   : ${RAG_CORPUS_DISPLAY_NAME}"
echo ""

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

gcloud services enable aiplatform.googleapis.com \
  storage.googleapis.com \
  --project="${PROJECT_ID}" --quiet

echo "✓ APIs enabled"

# ── 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 ─────────────────────────────────────
SEED_DIR="$(dirname "$0")/../docs/corpus-seed"
if [[ -d "${SEED_DIR}" ]] && ls "${SEED_DIR}"/*.md &>/dev/null; then
  gsutil -m cp "${SEED_DIR}"/*.md "gs://${CORPUS_BUCKET}/seed/" 2>/dev/null || true
  echo "✓ Seed documents uploaded to gs://${CORPUS_BUCKET}/seed/"
else
  echo "  No seed documents found in docs/corpus-seed/ — skipping upload"
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.")
    print("Run: pip install google-cloud-aiplatform>=1.87.0")
    sys.exit(1)

PROJECT_ID        = os.environ["PROJECT_ID"]
RAG_REGION        = os.environ.get("RAG_REGION", os.environ["REGION"])
FALLBACK_REGION   = "us-east1"
DISPLAY_NAME      = os.environ["RAG_CORPUS_DISPLAY_NAME"]
CORPUS_BUCKET     = f"{PROJECT_ID}-agent-corpus"

def try_init_and_list(region):
    vertexai.init(project=PROJECT_ID, location=region)
    try:
        return list(rag.list_corpora()), region
    except Exception as e:
        if "PERMISSION_DENIED" in str(e) or "allowlist" in str(e).lower() or "not found" in str(e).lower():
            return None, region
        raise

# Try primary region, fall back if needed
corpora, active_region = try_init_and_list(RAG_REGION)
if corpora is None and RAG_REGION != FALLBACK_REGION:
    print(f"  RAG Engine not available in {RAG_REGION} — falling back to {FALLBACK_REGION}")
    corpora, active_region = try_init_and_list(FALLBACK_REGION)

if corpora is None:
    print(f"ERROR: RAG Engine not accessible in {RAG_REGION} or {FALLBACK_REGION}.")
    print("Apply for allowlist: vertex-ai-rag-engine-support@google.com")
    print("Or set RAG_REGION=us-east1 in .env")
    sys.exit(1)

print(f"✓ RAG Engine accessible in region: {active_region}")

# Find or create corpus
corpus = None
for c in corpora:
    if c.display_name == 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=DISPLAY_NAME,
        embedding_model_config=embedding_config,
    )
    print(f"✓ RAG corpus created: {corpus.name}")

# Import seed documents (non-fatal)
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(f"  Import manually: https://console.cloud.google.com/vertex-ai/rag?project={PROJECT_ID}")

# Persist corpus name for 08-memorybank-setup.sh and agent
with open("/tmp/rag_corpus_name.txt", "w") as f:
    f.write(corpus.name)

print(f"")
print(f"  Corpus resource name : {corpus.name}")
print(f"  Active region        : {active_region}")
print(f"")
print(f"  ACTION REQUIRED — add to .env:")
print(f"    RAG_CORPUS_NAME={corpus.name}")
if active_region != os.environ.get("REGION"):
    print(f"    RAG_REGION={active_region}")
PYEOF

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