#!/usr/bin/env bash
# 07-rag-setup.sh — Create Vertex AI RAG Engine corpus in Serverless mode
# Uses REST API directly — no SDK version dependency
# RAG_REGION defaults to europe-west4 (Serverless available, no allowlist)
# Idempotent — safe to run multiple times
# 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}"

# europe-west4: Serverless RAG available, no allowlist, gemini-2.0-flash + 2.5-pro present
RAG_REGION="${RAG_REGION:-europe-west4}"

echo "=== 07: Setting up Vertex AI RAG Engine (Serverless mode) ==="
echo "    Project  : ${PROJECT_ID}"
echo "    Region   : ${RAG_REGION}  (override with RAG_REGION= if needed)"
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 (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 docs ──────────────────────────────────────────────────────────
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 in docs/corpus-seed/ — skipping"
fi

# ── REST: check if corpus exists ───────────────────────────────────────────
TOKEN=$(gcloud auth print-access-token)
API="https://${RAG_REGION}-aiplatform.googleapis.com/v1beta1"
PARENT="projects/${PROJECT_ID}/locations/${RAG_REGION}"

echo "  Checking for existing corpus..."
EXISTING=$(curl -sf -H "Authorization: Bearer ${TOKEN}" \
  "${API}/${PARENT}/ragCorpora" 2>/dev/null || echo '{}')

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

if [[ -n "${CORPUS_NAME}" ]]; then
  echo "✓ RAG corpus already exists: ${CORPUS_NAME}"
else
  echo "  Creating corpus '${RAG_CORPUS_DISPLAY_NAME}' in ${RAG_REGION} (Serverless/REST)..."

  RESPONSE=$(curl -sf -X POST \
    -H "Authorization: Bearer ${TOKEN}" \
    -H "Content-Type: application/json" \
    "${API}/${PARENT}/ragCorpora" \
    -d '{
      "displayName": "'"${RAG_CORPUS_DISPLAY_NAME}"'",
      "ragEmbeddingModelConfig": {
        "vertexPredictionEndpoint": {
          "publisherModel": "publishers/google/models/text-embedding-005"
        }
      },
      "ragVectorDbConfig": {
        "ragManagedDb": {}
      }
    }')

  # REST create returns an LRO — poll until done
  OPERATION=$(echo "${RESPONSE}" | python3 -c "import sys,json; print(json.load(sys.stdin).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://${RAG_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" ]] || [[ "${DONE}" == "true" ]]; then
      CORPUS_NAME=$(echo "${LRO}" | python3 -c "
import sys, json
d = json.load(sys.stdin)
print(d.get('response', {}).get('name', '') or d.get('metadata', {}).get('genericMetadata', {}).get('updateTime', ''))
" 2>/dev/null || true)
      # Fallback: re-list to get name
      if [[ -z "${CORPUS_NAME}" ]] || [[ "${CORPUS_NAME}" != *ragCorpora* ]]; then
        CORPUS_NAME=$(curl -sf -H "Authorization: Bearer ${TOKEN}" \
          "${API}/${PARENT}/ragCorpora" 2>/dev/null \
          | python3 -c "
import sys, json
data = json.load(sys.stdin)
for c in data.get('ragCorpora', []):
    if c.get('displayName') == '${RAG_CORPUS_DISPLAY_NAME}':
        print(c['name'])
        break
" 2>/dev/null || true)
      fi
      echo "✓ RAG corpus created: ${CORPUS_NAME}"
      break
    fi
    echo "  ... still waiting (${i}/30)"
  done

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

# ── Import seed docs via REST ─────────────────────────────────────────────
GCS_URI="gs://${CORPUS_BUCKET}/seed/"
curl -sf -X POST \
  -H "Authorization: Bearer ${TOKEN}" \
  -H "Content-Type: application/json" \
  "${API}/${CORPUS_NAME}:importRagFiles" \
  -d '{
    "importRagFilesConfig": {
      "gcsSource": { "uris": ["'"${GCS_URI}"'"] },
      "ragFileChunkingConfig": { "chunkSize": 512, "chunkOverlap": 50 }
    }
  }' &>/dev/null && echo "✓ Seed import triggered (async)" || echo "  WARNING: Seed import skipped (non-fatal)"

# ── Write output ──────────────────────────────────────────────────────────────────
echo "${CORPUS_NAME}" > /tmp/rag_corpus_name.txt

echo ""
echo "  Corpus resource name : ${CORPUS_NAME}"
echo "  Region               : ${RAG_REGION}"
echo ""
echo "  ACTION REQUIRED — add to .env:"
echo "    export RAG_CORPUS_NAME=\"${CORPUS_NAME}\""
echo "    export RAG_REGION=\"${RAG_REGION}\""
echo ""
echo "=== 07: RAG Engine setup COMPLETE ==="
echo "    View: https://console.cloud.google.com/vertex-ai/rag?project=${PROJECT_ID}"
echo ""