#!/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 payload omits vectorDbConfig entirely. Sending
  vectorDbConfig.ragManagedDb forces Spanner mode which is allowlist-restricted
  for new projects. Omitting it honours the project-level basic (serverless) config.
"""

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 ragEngineConfig endpoint is global/control-plane
    and does not exist in other regions (patching europe-west1 etc. returns 'Invalid
    endpoint name').
    """
    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"
    )
    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"  ⚠️  Engine config warning ({err.get('code')}): {err.get('message')} — 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("  ⚠️  Operation timed out — continuing anyway.")
            return

    print("  ✓ 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.

    IMPORTANT: vectorDbConfig is intentionally omitted from the payload.
    Sending vectorDbConfig.ragManagedDb explicitly requests Spanner mode,
    which is allowlist-restricted for new projects. Omitting it causes the
    API to honour the project-level ragEngineConfig (basic/serverless).
    """
    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"
            }
        }
        # vectorDbConfig intentionally omitted — see docstring above.
    })
    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"✓ 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("✓ 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"✓ 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()