"""
setup_indexes.py
================
PoC Task 2 - Neo4j Vector Migration
Task 04: One-time initialization script untuk membuat 5 Vector Index di Neo4j AuraDB.

Jalankan SEKALI saja sebelum mulai ingesting data:
  python setup_indexes.py

Vector Index yang dibuat (1024 dimensi, Cosine Similarity):
  1. background_emb    -> embedding_background
  2. methodology_emb   -> embedding_methodology
  3. gap_emb           -> embedding_gap
  4. objective_emb     -> embedding_objective
  5. futurework_emb    -> embedding_futurework
"""

import os
import sys
from dotenv import load_dotenv
from neo4j import GraphDatabase

# --- Load credentials dari .env ---
load_dotenv()

NEO4J_URI      = os.getenv("NEO4J_URI", "")
NEO4J_USERNAME = os.getenv("NEO4J_USERNAME", "neo4j")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "")

# --- Validasi credentials ---
if not all([NEO4J_URI, NEO4J_USERNAME, NEO4J_PASSWORD]):
    print("ERROR: Kredensial Neo4j belum lengkap di file .env")
    print("   Pastikan NEO4J_URI, NEO4J_USERNAME, dan NEO4J_PASSWORD sudah terisi.")
    sys.exit(1)

# --- Daftar 5 Vector Index yang akan dibuat ---
# Format: (index_name, property_name)
VECTOR_INDEXES = [
    ("background_emb",   "embedding_background"),
    ("methodology_emb",  "embedding_methodology"),
    ("gap_emb",          "embedding_gap"),
    ("objective_emb",    "embedding_objective"),
    ("futurework_emb",   "embedding_futurework"),
]

DIMENSIONS = 1536
SIMILARITY_FUNCTION = "cosine"
NODE_LABEL = "Literature"


def create_vector_index(driver, index_name: str, property_name: str):
    """Membuat satu vector index pada Neo4j AuraDB (IF NOT EXISTS)."""
    cypher = f"""
    CREATE VECTOR INDEX {index_name} IF NOT EXISTS
    FOR (n:{NODE_LABEL}) ON (n.{property_name})
    OPTIONS {{
        indexConfig: {{
            `vector.dimensions`: {DIMENSIONS},
            `vector.similarity_function`: '{SIMILARITY_FUNCTION}'
        }}
    }}
    """
    with driver.session() as session:
        session.run(cypher)
    print(f" Index '{index_name}' -> (:{NODE_LABEL}.{property_name}) → CREATED / ALREADY EXISTS")


def verify_indexes(driver):
    """Menampilkan daftar vector index yang aktif di Neo4j."""
    print("\n Verifikasi Vector Indexes yang Aktif:")
    print("-" * 55)
    with driver.session() as session:
        result = session.run("SHOW VECTOR INDEXES YIELD name, state, labelsOrTypes, properties")
        rows = list(result)
        if not rows:
            print("Tidak ada vector index yang ditemukan !!!.")
            return
        for row in rows:
            name   = row["name"]
            state  = row["state"]
            labels = row["labelsOrTypes"]
            props  = row["properties"]
            status = "Success" if state == "ONLINE" else "Pending"
            print(f"  {status} [{state}] {name} → {labels}.{props}")
    print("-" * 55)


def main():
    print("=" * 55)
    print("  PoC Task 2 - Neo4j Vector Index Setup")
    print(f"  URI     : {NEO4J_URI[:30]}...")
    print(f"  User    : {NEO4J_USERNAME}")
    print(f"  Dimensi : {DIMENSIONS}D | Similarity: {SIMILARITY_FUNCTION.upper()}")
    print("=" * 55)

    driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))

    try:
        # Verifikasi koneksi
        driver.verify_connectivity()
        print("\nKoneksi ke Neo4j AuraDB: BERHASIL\n")

        # Buat semua vector indexes
        print("Membuat 5 Vector Indexes...")
        for index_name, property_name in VECTOR_INDEXES:
            create_vector_index(driver, index_name, property_name)

        # Verifikasi hasil
        verify_indexes(driver)

        print("\n Setup selesai Database siap menerima embedding 1536D.\n")

    except Exception as e:
        print(f"\n ERROR saat menghubungi Neo4j: {e}")
        sys.exit(1)
    finally:
        driver.close()


if __name__ == "__main__":
    main()
