#!/usr/bin/env python3
"""
setup_corpus.py — Create a Vertex AI RAG Engine corpus and import documents.
Project: propane-will-491900-m5

SDK 1.153.1 has a bug where backend_config crashes when provided,
and defaults to Spanner when omitted. We bypass it by calling the
REST API directly for corpus creation, then use the SDK for everything else.

Key design decisions:
- ragEngineConfig PATCH always targets us-central1 — it is a project-level
  control-plane endpoint that does NOT exist in other regions.
- Corpus creation uses ragManagedDbConfig.basic in the RagCorpus body.
  The v1beta1 API uses ragManagedDbConfig (not vectorDbConfig) on the corpus
  resource. Sending it explicitly with {"basic":{}} forces serverless mode
  and overrides any project-level Spanner default.
"""

import os
import json
import time
import subprocess

import vertexai
from vertexai import rag
from vertexai.rag import RagCorpus

PROJECT_ID = "propane-will-491900-m5"
LOCATION = os.environ.get("RAG_LOCATION", "us-central1")
CORPUS_DISPLAY_NAME = os.environ.get("RAG_CORPUS_NAME", "osvauco-knowledge-base")
GCS_SOURCE = os.environ.get(
    "RAG_GCS_SOURCE",
    f"gs://{PROJECT_ID}-agent-staging/rag-docs/",
)

# ragEngineConfig is a project-level control-plane endpoint.
# It only exists in us-central1 regardless of where the corpus lives.
_ENGINE_CONFIG_LOCATION = "us-central1"


def get_token() -> str:
    return subprocess.check_output(
        ["gcloud", "auth", "print-access-token"], text=True
    ).strip()


def ensure_serverless_engine_config() -> None:
    """
    Set project-level RAG Engine Config to basic (serverless) tier.
    Always targets us-central1. The endpoint does not exist in other regions.

    The v1beta1 ragEngineConfig schema:
      { "ragManagedDbConfig": { "basic": {} } }   <- serverless
      { "ragManagedDbConfig": { "scaled": {} } }  <- Spanner / managed DB
    """
    print("Ensuring RAG Engine Config is set to serverless (basic) tier...")
    endpoint = (
        f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1"
        f"/projects/{PROJECT_ID}/locations/{_ENGINE_CONFIG_LOCATION}/ragEngineConfig"
    )
    # Explicitly set basic and unset spanner by sending only basic.
    payload = json.dumps({"ragManagedDbConfig": {"basic": {}}})
    token = get_token()

    result = subprocess.run(
        ["curl", "-s", "-X", "PATCH",
         "-H", f"Authorization: Bearer {token}",
         "-H", "Content-Type: application/json",
         endpoint, "-d", payload],
        capture_output=True, text=True,
    )
    resp = json.loads(result.stdout)

    if "error" in resp:
        err = resp["error"]
        print(f"  \u26a0\ufe0f  Engine config warning ({err.get('code')}): {err.get('message')} \u2014 continuing.")
        return

    op_name = resp.get("name", "")
    if "/operations/" in op_name and not resp.get("done"):
        print(f"  Polling operation {op_name.split('/')[-1]}...")
        op_url = (
            f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1/{op_name}"
        )
        for _ in range(20):
            time.sleep(3)
            token = get_token()
            r = subprocess.run(
                ["curl", "-s", "-H", f"Authorization: Bearer {token}", op_url],
                capture_output=True, text=True,
            )
            op = json.loads(r.stdout)
            if op.get("done"):
                break
        else:
            print("  \u26a0\ufe0f  Operation timed out \u2014 continuing anyway.")
            return

    print("  \u2713 RAG Engine Config set to basic tier.")


def create_corpus_rest() -> str:
    """
    Create corpus via REST API directly, bypassing SDK backend_config bug.
    Returns the corpus resource name.

    Uses ragManagedDbConfig.basic explicitly in the corpus body.
    The v1beta1 RagCorpus resource accepts ragManagedDbConfig directly
    (not vectorDbConfig). Sending {"basic":{}} forces serverless mode
    regardless of the project-level engine config default.
    """
    url = (
        f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1"
        f"/projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora"
    )
    payload = json.dumps({
        "displayName": CORPUS_DISPLAY_NAME,
        "ragEmbeddingModelConfig": {
            "vertexPredictionEndpoint": {
                "model": "publishers/google/models/text-embedding-004"
            }
        },
        # Explicitly request serverless (basic) mode at the corpus level.
        # The v1beta1 API uses ragManagedDbConfig on the RagCorpus body,
        # NOT vectorDbConfig. This overrides any Spanner project default.
        "ragManagedDbConfig": {
            "basic": {}
        }
    })
    token = get_token()
    result = subprocess.run(
        ["curl", "-s", "-X", "POST",
         "-H", f"Authorization: Bearer {token}",
         "-H", "Content-Type: application/json",
         url, "-d", payload],
        capture_output=True, text=True,
    )
    resp = json.loads(result.stdout)
    if "error" in resp:
        raise RuntimeError(f"Failed to create corpus: {resp['error']}")

    # REST returns a long-running operation
    op_name = resp.get("name", "")
    if "/operations/" not in op_name:
        raise RuntimeError(f"Unexpected response: {resp}")

    print(f"  Polling corpus creation operation...")
    op_url = f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{op_name}"
    for _ in range(40):
        time.sleep(5)
        token = get_token()
        r = subprocess.run(
            ["curl", "-s", "-H", f"Authorization: Bearer {token}", op_url],
            capture_output=True, text=True,
        )
        op = json.loads(r.stdout)
        if op.get("done"):
            if "error" in op:
                raise RuntimeError(f"Corpus creation failed: {op['error']}")
            corpus_name = op["response"]["name"]
            return corpus_name

    raise RuntimeError("Corpus creation operation timed out.")


def get_or_create_corpus() -> RagCorpus:
    """Return existing corpus by display name, or create a new serverless one."""
    for c in rag.list_corpora():
        if c.display_name == CORPUS_DISPLAY_NAME:
            print(f"Corpus '{CORPUS_DISPLAY_NAME}' already exists: {c.name}")
            return c

    print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} via REST...")
    corpus_name = create_corpus_rest()
    print(f"\u2713 Corpus created: {corpus_name}")

    # Re-fetch via SDK so we get a proper RagCorpus object
    for c in rag.list_corpora():
        if c.name == corpus_name:
            return c

    raise RuntimeError(f"Corpus created but not found in list: {corpus_name}")


def import_documents(corpus: RagCorpus) -> None:
    print(f"Importing files from {GCS_SOURCE}...")
    rag.import_files(
        corpus.name,
        paths=[GCS_SOURCE],
        transformation_config=rag.TransformationConfig(
            chunking_config=rag.ChunkingConfig(
                chunk_size=512,
                chunk_overlap=100,
            )
        ),
    )
    print("\u2713 Import complete.")


def test_retrieval(corpus: RagCorpus) -> None:
    print("Running test retrieval query...")
    response = rag.retrieval_query(
        rag_resources=[rag.RagResource(rag_corpus=corpus.name)],
        text="test query",
        rag_retrieval_config=rag.RagRetrievalConfig(top_k=3),
    )
    print(f"\u2713 Test retrieval returned {len(response.contexts.contexts)} chunk(s).")


def main() -> None:
    vertexai.init(project=PROJECT_ID, location=LOCATION)

    ensure_serverless_engine_config()

    corpus = get_or_create_corpus()

    import_documents(corpus)

    print(f"\nRAG_CORPUS={corpus.name}")
    print("Add this to Secret Manager:")
    print(f"  gcloud secrets create rag-corpus-name --data-file=- <<<'{corpus.name}'")

    test_retrieval(corpus)


if __name__ == "__main__":
    main()