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

IMPORTANT: RAG Engine in us-central1 requires allowlist access.
Contact: vertex-ai-rag-engine-support@google.com
Alternative: Use us-east1 or us-east4 for immediate access.
"""

import os
import vertexai
from vertexai import rag

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

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

    # Check if corpus already exists
    existing = list(rag.list_corpora())
    for c in existing:
        if c.display_name == CORPUS_DISPLAY_NAME:
            print(f"Corpus '{CORPUS_DISPLAY_NAME}' already exists: {c.name}")
            corpus = c
            break
    else:
        print(f"Creating RAG corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION}...")
        corpus = rag.create_corpus(
            display_name=CORPUS_DISPLAY_NAME,
        )
        print(f"Corpus created: {corpus.name}")

    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.")
    print(f"\nRAG_CORPUS={corpus.name}")
    print("Set this as an environment variable or Secret Manager entry.")

    # Test retrieval
    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)} chunks.")


if __name__ == "__main__":
    main()