#!/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)
# Uploads ALL docs/*.md + corpus-seed/*.md + protocols/sessions/*.md to GCS
# 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:-europe-west4}"
REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"

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"

# ── Hent project-number (kræves av importRagFiles-endepunktet) ────────────────
PROJECT_NUMBER=$(gcloud projects describe "${PROJECT_ID}" \
  --format='value(projectNumber)' 2>/dev/null)
echo "  Project number: ${PROJECT_NUMBER}"

# ── 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 alle kildedokumenter til GCS ──────────────────────────────────────
UPLOAD_COUNT=0

SEED_DIR="${REPO_ROOT}/docs/corpus-seed"
if [[ -d "${SEED_DIR}" ]] && compgen -G "${SEED_DIR}/*.md" > /dev/null 2>&1; then
  gsutil -m cp "${SEED_DIR}"/*.md "gs://${CORPUS_BUCKET}/seed/" 2>/dev/null || true
  COUNT=$(ls "${SEED_DIR}"/*.md 2>/dev/null | wc -l | tr -d ' ')
  echo "✓ corpus-seed/: ${COUNT} filer → gs://${CORPUS_BUCKET}/seed/"
  UPLOAD_COUNT=$((UPLOAD_COUNT + COUNT))
fi

DOCS_DIR="${REPO_ROOT}/docs"
if compgen -G "${DOCS_DIR}/*.md" > /dev/null 2>&1; then
  gsutil -m cp "${DOCS_DIR}"/*.md "gs://${CORPUS_BUCKET}/docs/" 2>/dev/null || true
  COUNT=$(ls "${DOCS_DIR}"/*.md 2>/dev/null | wc -l | tr -d ' ')
  echo "✓ docs/: ${COUNT} filer → gs://${CORPUS_BUCKET}/docs/"
  UPLOAD_COUNT=$((UPLOAD_COUNT + COUNT))
fi

SESSIONS_DIR="${REPO_ROOT}/protocols/sessions"
if [[ -d "${SESSIONS_DIR}" ]] && compgen -G "${SESSIONS_DIR}/*.md" > /dev/null 2>&1; then
  gsutil -m cp "${SESSIONS_DIR}"/*.md "gs://${CORPUS_BUCKET}/sessions/" 2>/dev/null || true
  COUNT=$(ls "${SESSIONS_DIR}"/*.md 2>/dev/null | wc -l | tr -d ' ')
  echo "✓ protocols/sessions/: ${COUNT} filer → gs://${CORPUS_BUCKET}/sessions/"
  UPLOAD_COUNT=$((UPLOAD_COUNT + COUNT))
fi

if [[ -f "${REPO_ROOT}/MASTERPLAN.md" ]]; then
  gsutil cp "${REPO_ROOT}/MASTERPLAN.md" "gs://${CORPUS_BUCKET}/MASTERPLAN.md" 2>/dev/null || true
  echo "✓ MASTERPLAN.md → gs://${CORPUS_BUCKET}/MASTERPLAN.md"
  UPLOAD_COUNT=$((UPLOAD_COUNT + 1))
fi
echo "  Totalt ${UPLOAD_COUNT} filer synkronisert"

# ── REST: hent/opprett corpus (bruker project-ID for listing/oppretting) ────────
TOKEN=$(gcloud auth print-access-token)
BASE_URL="https://${RAG_REGION}-aiplatform.googleapis.com/v1beta1"
PARENT="projects/${PROJECT_ID}/locations/${RAG_REGION}"

echo "  Sjekker om corpus finnes..."
EXISTING=$(curl -sf -H "Authorization: Bearer ${TOKEN}" \
  "${BASE_URL}/${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 finnes allerede: ${CORPUS_NAME}"
else
  echo "  Oppretter corpus '${RAG_CORPUS_DISPLAY_NAME}' i ${RAG_REGION}..."
  RESPONSE=$(curl -sf -X POST \
    -H "Authorization: Bearer ${TOKEN}" \
    -H "Content-Type: application/json" \
    "${BASE_URL}/${PARENT}/ragCorpora" \
    -d '{"displayName":"'"${RAG_CORPUS_DISPLAY_NAME}"'","ragEmbeddingModelConfig":{"vertexPredictionEndpoint":{"publisherModel":"publishers/google/models/text-embedding-005"}},"ragVectorDbConfig":{"ragManagedDb":{}}}')
  OPERATION=$(echo "${RESPONSE}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('name',''))" 2>/dev/null || true)
  [[ -z "${OPERATION}" ]] && { echo "ERROR: ${RESPONSE}"; exit 1; }
  echo "  Venter på LRO..."
  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" ]]; then
      CORPUS_NAME=$(curl -sf -H "Authorization: Bearer ${TOKEN}" \
        "${BASE_URL}/${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)
      echo "✓ Corpus opprettet: ${CORPUS_NAME}"
      break
    fi
    echo "  ... venter (${i}/30)"
  done
  [[ -z "${CORPUS_NAME}" ]] && { echo "ERROR: timed out"; exit 1; }
fi

# Hent corpus-ID (siste del av ressursstien)
CORPUS_ID=$(echo "${CORPUS_NAME}" | rev | cut -d'/' -f1 | rev)

# ── Import: bygg URL med project-NUMBER (ikke project-ID) ───────────────────────
# importRagFiles krever project-number i stien
IMPORT_URL="https://${RAG_REGION}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_NUMBER}/locations/${RAG_REGION}/ragCorpora/${CORPUS_ID}:importRagFiles"
echo "  Import URL: ${IMPORT_URL}"

IMPORT_OK=0
for GCS_PATH in "seed/" "docs/" "sessions/" "MASTERPLAN.md"; do
  GCS_URI="gs://${CORPUS_BUCKET}/${GCS_PATH}"
  if gsutil ls "${GCS_URI}" &>/dev/null; then
    HTTP_STATUS=$(curl -s -o /tmp/import_response.json -w "%{http_code}" -X POST \
      -H "Authorization: Bearer ${TOKEN}" \
      -H "Content-Type: application/json" \
      "${IMPORT_URL}" \
      -d '{"importRagFilesConfig":{"gcsSource":{"uris":["'"${GCS_URI}"'"]},"ragFileChunkingConfig":{"chunkSize":512,"chunkOverlap":50}}}')
    if [[ "${HTTP_STATUS}" == "200" ]]; then
      echo "✓ Import trigget: ${GCS_URI}"
      IMPORT_OK=$((IMPORT_OK + 1))
    else
      echo "  ADVARSEL: Import feilet (HTTP ${HTTP_STATUS}): ${GCS_URI}"
      cat /tmp/import_response.json 2>/dev/null || true
    fi
  fi
done
echo "  ${IMPORT_OK}/4 import-jobber trigget (async — indeksering tar noen minutter)"

# ── Output ─────────────────────────────────────────────────────────────────
echo "${CORPUS_NAME}" > /tmp/rag_corpus_name.txt
echo ""
echo "  Corpus resource name : ${CORPUS_NAME}"
echo "  Region               : ${RAG_REGION}"
echo ""
echo "  ACTION REQUIRED — legg til i .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 ""