#!/usr/bin/env bash
# 07-rag-setup.sh — Create Vertex AI RAG Engine corpus in Serverless mode
# Serverless = RagManagedDb (no Spanner, no allowlist needed)
# 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}"

RAG_REGION="${RAG_REGION:-${REGION}}"

echo "=== 07: Setting up Vertex AI RAG Engine (Serverless mode) ==="
echo "    Project  : ${PROJECT_ID}"
echo "    Region   : ${RAG_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"

# ── Upgrade SDK to ensure RagManagedDb support ───────────────────────────────
echo "  Upgrading google-cloud-aiplatform SDK..."
pip install --quiet --upgrade google-cloud-aiplatform
echo "✓ SDK upgraded"

# ── 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

# ── Python: create corpus ────────────────────────────────────────────────────
python3 - << 'PYEOF'
import os, sys

try:
    import vertexai
    from vertexai.preview import rag
except ImportError:
    print("ERROR: google-cloud-aiplatform not installed.")
    sys.exit(1)

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

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

# Check if corpus already exists
corpus = None
try:
    for c in rag.list_corpora():
        if c.display_name == DISPLAY_NAME:
            corpus = c
            print(f"✓ RAG corpus already exists: {c.name}")
            break
except Exception as e:
    print(f"WARNING: Could not list corpora: {e}")

if corpus is None:
    print(f"  Creating corpus '{DISPLAY_NAME}' in {RAG_REGION}...")
    embedding_config = rag.EmbeddingModelConfig(
        publisher_model="publishers/google/models/text-embedding-005"
    )

    # Attempt order:
    # 1. SDK >= 1.87 : RagVectorDbConfig(rag_managed_db=RagManagedDb())
    # 2. Older SDK   : RagVectorDbConfig(rag_managed_db=RagManagedDbConfig())
    # 3. REST fallback: gapic-style with vector_db proto dict
    attempts = ["new_sdk", "old_sdk", "proto_dict"]
    for attempt in attempts:
        try:
            if attempt == "new_sdk":
                vector_db = rag.RagVectorDbConfig(
                    rag_managed_db=rag.RagManagedDb()
                )
                corpus = rag.create_corpus(
                    display_name=DISPLAY_NAME,
                    embedding_model_config=embedding_config,
                    vector_db=vector_db,
                )

            elif attempt == "old_sdk":
                # Some SDK versions use RagManagedDbConfig instead
                RagManagedDbConfig = getattr(rag, "RagManagedDbConfig", None)
                if RagManagedDbConfig is None:
                    raise AttributeError("RagManagedDbConfig not in this SDK version")
                vector_db = rag.RagVectorDbConfig(
                    rag_managed_db=RagManagedDbConfig()
                )
                corpus = rag.create_corpus(
                    display_name=DISPLAY_NAME,
                    embedding_model_config=embedding_config,
                    vector_db=vector_db,
                )

            elif attempt == "proto_dict":
                # Last resort: pass vector_db as a plain dict understood by gapic
                from google.cloud.aiplatform_v1beta1.types import (
                    RagCorpus, RagVectorDbConfig, RagManagedDb,
                    RagEmbeddingModelConfig
                )
                from google.cloud import aiplatform_v1beta1 as aip

                client = aip.VertexRagDataServiceClient(
                    client_options={"api_endpoint": f"{RAG_REGION}-aiplatform.googleapis.com"}
                )
                parent = f"projects/{PROJECT_ID}/locations/{RAG_REGION}"
                rag_corpus = RagCorpus(
                    display_name=DISPLAY_NAME,
                    rag_embedding_model_config=RagEmbeddingModelConfig(
                        vertex_prediction_endpoint=RagEmbeddingModelConfig.VertexPredictionEndpoint(
                            publisher_model="publishers/google/models/text-embedding-005"
                        )
                    ),
                    rag_vector_db_config=RagVectorDbConfig(
                        rag_managed_db=RagManagedDb()
                    ),
                )
                op = client.create_rag_corpus(parent=parent, rag_corpus=rag_corpus)
                result = op.result()
                # Wrap in a simple namespace so rest of script works
                class _C:
                    name = result.name
                corpus = _C()

            print(f"✓ RAG corpus created [{attempt}]: {corpus.name}")
            break

        except Exception as e:
            if attempt == attempts[-1]:
                print(f"ERROR: All attempts failed. Last error: {e}")
                sys.exit(1)
            print(f"  [{attempt}] failed: {e} — trying next method...")

# 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: {e}")

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"  Region               : {RAG_REGION}")
print(f"")
print(f"  ACTION REQUIRED — add to .env:")
print(f'    export RAG_CORPUS_NAME="{corpus.name}"')
PYEOF

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