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

Explicitly uses Serverless (RagManagedDbConfig basic tier) to avoid
Spanner allowlist restrictions on new projects in us-central1.
"""

import os
import json
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/",
)


def ensure_serverless_engine_config() -> None:
    """
    Pre-flight: set project-level RAG Engine Config to basic (serverless) tier.
    Idempotent — safe to re-run.
    """
    print("Ensuring RAG Engine Config is set to serverless (basic) tier...")
    endpoint = (
        f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1"
        f"/projects/{PROJECT_ID}/locations/{LOCATION}/ragEngineConfig"
    )
    payload = json.dumps({"ragManagedDbConfig": {"basic": {}}})

    try:
        token = subprocess.check_output(
            ["gcloud", "auth", "print-access-token"], text=True
        ).strip()
    except subprocess.CalledProcessError:
        print("  ⚠️  Could not get gcloud token — skipping pre-flight.")
        return

    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 = result.stdout
    if '"error"' in resp:
        err = json.loads(resp).get("error", {})
        code, msg = err.get("code"), err.get("message", "")
        if code == 400 and "already" in msg.lower():
            print(f"  ✓ Already configured: {msg}")
        else:
            print(f"  ⚠️  Engine config warning ({code}): {msg} — continuing anyway.")
    else:
        print("  ✓ RAG Engine Config set to basic tier.")


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} (serverless)...")
    corpus = rag.create_corpus(
        display_name=CORPUS_DISPLAY_NAME,
        # Explicitly request serverless / basic managed DB tier
        backend_config=rag.RagVectorDbConfig(
            rag_managed_db=rag.RagManagedDbConfig(
                tier=rag.RagManagedDbConfig.Tier.BASIC,
            )
        ),
    )
    print(f"✓ Corpus created: {corpus.name}")
    return corpus


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()