#!/usr/bin/env python3
"""
sync_corpus.py — Synchronize RAG Engine corpus with GCS and perform health checks.
Project: propane-will-491900-m5

Capabilities:
- Locates existing corpus by display name.
- Lists current files in the corpus.
- Triggers incremental import from GCS.
- Performs a health-check retrieval query.
- Reports status and chunk counts.

Graceful degradation:
- If corpus is not found (e.g. due to Spanner restriction), logs a warning and exits cleanly.
"""

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

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

def find_corpus() -> RagCorpus:
    """Search for the corpus in the configured location and fallback location."""
    locations_to_try = [LOCATION]
    if LOCATION != "europe-west4":
        locations_to_try.append("europe-west4")
    
    for loc in locations_to_try:
        print(f"Searching for corpus '{CORPUS_DISPLAY_NAME}' in {loc}...")
        vertexai.init(project=PROJECT_ID, location=loc)
        try:
            for c in rag.list_corpora():
                if c.display_name == CORPUS_DISPLAY_NAME:
                    print(f"  \u2713 Found: {c.name}")
                    return c
        except Exception as e:
            print(f"  \u26a0\ufe0f  Notice: Could not list corpora in {loc}: {e}")
            
    return None

def list_files(corpus: RagCorpus):
    """List files currently in the corpus."""
    print(f"Listing files in corpus...")
    files = list(rag.list_files(corpus_name=corpus.name))
    if not files:
        print("  (Corpus is empty)")
        return []
    
    for f in files:
        print(f"  - {f.display_name} (state: {f.state.name})")
    return files

def sync_from_gcs(corpus: RagCorpus):
    """Trigger incremental import from GCS."""
    print(f"Syncing from {GCS_SOURCE}...")
    # import_files in SDK is essentially an incremental trigger
    rag.import_files(
        corpus.name,
        paths=[GCS_SOURCE],
        transformation_config=rag.TransformationConfig(
            chunking_config=rag.ChunkingConfig(
                chunk_size=512,
                chunk_overlap=100,
            )
        ),
    )
    print("  \u2713 Sync job submitted.")

def run_health_check(corpus: RagCorpus):
    """Perform a retrieval query and report results."""
    print("Running health check retrieval...")
    try:
        response = rag.retrieval_query(
            rag_resources=[rag.RagResource(rag_corpus=corpus.name)],
            text="How does OSVauco handle operator experience?",
            rag_retrieval_config=rag.RagRetrievalConfig(top_k=3),
        )
        contexts = response.contexts.contexts
        print(f"  \u2713 Health check returned {len(contexts)} contexts.")
        for i, ctx in enumerate(contexts):
            print(f"    [{i+1}] Source: {ctx.source_uri} (Score: {ctx.score:.4f})")
    except Exception as e:
        print(f"  \u2718 Health check failed: {e}")

def main():
    print(f"--- RAG Sync & Maintenance Tool ---")
    
    corpus = find_corpus()
    if not corpus:
        print(f"\n\u26a0\ufe0f  WARN: Corpus '{CORPUS_DISPLAY_NAME}' not found.")
        print("    RAG Engine might be restricted for this project or setup_corpus.py has not been run.")
        print("    Skipping sync and health check.")
        sys.exit(0)
    
    # Update global LOCATION if we found it in a different one
    found_loc = corpus.name.split('/')[3]
    print(f"Active Location: {found_loc}")
    
    list_files(corpus)
    
    # Always trigger a sync to ensure new files are picked up
    sync_from_gcs(corpus)
    
    print("\nWaiting 5s for sync initialization...")
    time.sleep(5)
    
    run_health_check(corpus)
    
    print(f"\nCOMPLETED")
    print(f"RAG_LOCATION={found_loc}")
    print(f"RAG_CORPUS={corpus.name}")

if __name__ == "__main__":
    main()