#!/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 bugs in backend_config and RagManagedDbConfig.
We use the REST API (v1beta1) for corpus creation and engine configuration.

Changes implemented:
- Graceful degradation: If RAG Engine is restricted (Spanner mode), skip and exit cleanly.
- ensure_serverless_engine_config() sets RAG_ENGINE_AVAILABLE flag.
- get_or_create_corpus() returns None if RAG is unavailable.
"""

import os
import json
import time
import subprocess
import sys

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

# Global status flag
RAG_ENGINE_AVAILABLE = True

# 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 _engine_config_url() -> str:
    return (
        f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1"
        f"/projects/{PROJECT_ID}/locations/{_ENGINE_CONFIG_LOCATION}/ragEngineConfig"
    )


def get_engine_config() -> dict:
    token = get_token()
    r = subprocess.run(
        ["curl", "-s", "-H", f"Authorization: Bearer {token}", _engine_config_url()],
        capture_output=True, text=True,
    )
    try:
        return json.loads(r.stdout)
    except Exception:
        print(f"Failed to parse engine config: {r.stdout}")
        return {}


def ensure_serverless_engine_config() -> None:
    """
    Set project-level RAG Engine Config to basic (serverless) tier.
    Only applicable for us-central1, us-east1, us-east4.
    """
    global RAG_ENGINE_AVAILABLE
    
    restricted_regions = {"us-central1", "us-east1", "us-east4"}
    if LOCATION not in restricted_regions:
        print(f"Skipping engine config check for location: {LOCATION} (not restricted)")
        return

    print(f"Ensuring RAG Engine Config in {_ENGINE_CONFIG_LOCATION} is set to serverless (basic) tier...")
    token = get_token()

    # Use updateMask=ragManagedDbConfig to target the entire configuration block.
    # Set spanner: null to explicitly clear it.
    payload = {
        "ragManagedDbConfig": {
            "basic": {},
            "spanner": None
        }
    }
    
    url = f"{_engine_config_url()}"
    
    result = subprocess.run(
        ["curl", "-s", "-X", "PATCH",
         "-H", f"Authorization: Bearer {token}",
         "-H", "Content-Type: application/json",
         url,
         "-d", json.dumps(payload)],
        capture_output=True, text=True,
    )
    
    try:
        resp = json.loads(result.stdout)
    except Exception:
        print(f"  \u26a0\ufe0f  Failed to parse PATCH response: {result.stdout}")
        return

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

    op_name = resp.get("name", "")
    if "/operations/" in op_name and not resp.get("done"):
        print(f"  Polling engine config operation...")
        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  Engine config operation timed out.")

    # Verify actual state
    cfg = get_engine_config()
    db_cfg = cfg.get("ragManagedDbConfig", {})
    print(f"  Current ragManagedDbConfig: {json.dumps(db_cfg)}")
    
    if "spanner" in db_cfg:
        print(f"\n\u26a0\ufe0f  WARN: Engine config in {_ENGINE_CONFIG_LOCATION} is still in Spanner mode.")
        print("    RAG Engine appears to be restricted for this project.")
        RAG_ENGINE_AVAILABLE = False
        return
            
    print("  \u2713 RAG Engine Config verified as basic tier.")


def create_corpus_rest(loc: str) -> str:
    """
    Create corpus via REST API directly.
    """
    url = (
        f"https://{loc}-aiplatform.googleapis.com/v1beta1"
        f"/projects/{PROJECT_ID}/locations/{loc}/ragCorpora"
    )
    # vectorDbConfig.ragManagedDb = serverless RAG Managed DB.
    # This is the correct v1beta1 field name.
    payload = json.dumps({
        "displayName": CORPUS_DISPLAY_NAME,
        "ragEmbeddingModelConfig": {
            "vertexPredictionEndpoint": {
                "model": "publishers/google/models/text-embedding-004"
            }
        },
        "vectorDbConfig": {
            "ragManagedDb": {}
        }
    })
    token = get_token()
    print(f"  Sending POST to {url}...")
    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,
    )
    
    try:
        resp = json.loads(result.stdout)
    except Exception:
        raise RuntimeError(f"Failed to parse API response: {result.stdout}")

    if "error" in resp:
        # Pass the error object up so get_or_create_corpus can inspect it
        raise RuntimeError(json.dumps(resp["error"]))

    op_name = resp.get("name", "")
    if "/operations/" not in op_name:
        raise RuntimeError(f"Unexpected response (no operation): {resp}")

    print(f"  Polling corpus creation operation {op_name}...")
    op_url = f"https://{loc}-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(json.dumps(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 with fallback."""
    global LOCATION, RAG_ENGINE_AVAILABLE
    
    if not RAG_ENGINE_AVAILABLE:
        print("\n\u2139  RAG Engine appears to be in restricted Spanner mode for this project. Skipping corpus creation.")
        return None

    def _find_in_list():
        try:
            for c in rag.list_corpora():
                if c.display_name == CORPUS_DISPLAY_NAME:
                    return c
        except Exception:
            pass
        return None

    # 1. Try finding in current LOCATION
    vertexai.init(project=PROJECT_ID, location=LOCATION)
    corpus = _find_in_list()
    if corpus:
        print(f"Corpus '{CORPUS_DISPLAY_NAME}' already exists in {LOCATION}: {corpus.name}")
        return corpus

    # 2. Try creating in current LOCATION
    print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} via REST...")
    try:
        corpus_name = create_corpus_rest(LOCATION)
        print(f"\u2713 Corpus created in {LOCATION}: {corpus_name}")
    except RuntimeError as e:
        err_str = str(e)
        # Check for Spanner restriction error message
        if "using Spanner mode with RAG Engine in us-central1, us-east1, and us-east4 is restricted" in err_str:
            if LOCATION != "europe-west4":
                print(f"\u26a0\ufe0f Spanner restriction detected in {LOCATION}. Retrying in europe-west4...")
                LOCATION = "europe-west4"
                vertexai.init(project=PROJECT_ID, location=LOCATION)
                
                # Check if it already exists in the fallback location
                corpus = _find_in_list()
                if corpus:
                    print(f"Corpus '{CORPUS_DISPLAY_NAME}' found in fallback {LOCATION}: {corpus.name}")
                    return corpus
                    
                try:
                    corpus_name = create_corpus_rest(LOCATION)
                    print(f"\u2713 Corpus created in fallback {LOCATION}: {corpus_name}")
                except RuntimeError as e2:
                    if "using Spanner mode" in str(e2):
                        print("\n\u26a0\ufe0f  RAG Engine restriction confirmed in fallback region. RAG is effectively unavailable.")
                        RAG_ENGINE_AVAILABLE = False
                        return None
                    raise
            else:
                print("\n\u26a0\ufe0f  RAG Engine restriction confirmed in fallback region. RAG is effectively unavailable.")
                RAG_ENGINE_AVAILABLE = False
                return None
        else:
            raise RuntimeError(f"Failed to create corpus in {LOCATION}: {err_str}")

    # Final retrieval of the corpus 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} into {corpus.name}...")
    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 job submitted.")


def test_retrieval(corpus: RagCorpus) -> None:
    print("Running test retrieval query...")
    try:
        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).")
    except Exception as e:
        print(f"  \u26a0\ufe0f  Retrieval test failed (import might still be processing): {e}")


def main() -> None:
    print(f"--- RAG Setup Starting (Target Location: {LOCATION}) ---")
    
    ensure_serverless_engine_config()

    corpus = get_or_create_corpus()

    if corpus is None:
        print("\n--- RAG AVAILABILITY SUMMARY ---")
        print(f"Project   : {PROJECT_ID}")
        print("Status    : UNAVAILABLE (Platform restriction: Spanner Mode only)")
        print("Action    : Skipping RAG integration. Downstream tools will operate without a corpus.")
        print("---------------------------------")
        sys.exit(0)

    import_documents(corpus)

    print(f"\nSUCCESS")
    print(f"RAG_LOCATION={LOCATION}")
    print(f"RAG_CORPUS={corpus.name}")
    print("\nNext steps (set in your environment):")
    print(f"export RAG_LOCATION={LOCATION}")
    print(f"export RAG_CORPUS_NAME={corpus.name}")
    
    print("\nWaiting 10s before test retrieval...")
    time.sleep(10)
    test_retrieval(corpus)


if __name__ == "__main__":
    main()