import os
import json
from dotenv import load_dotenv

load_dotenv()  # take environment variables from .env.

from cassandra.auth import PlainTextAuthProvider
from cassandra.cluster import Cluster
from llama_index import ServiceContext
from llama_index import set_global_service_context
from llama_index import VectorStoreIndex, SimpleDirectoryReader, StorageContext
from llama_index.embeddings import GradientEmbedding
from llama_index.llms import GradientBaseModelLLM
from llama_index.vector_stores import CassandraVectorStore

import cassandra
print("cassandra version:", cassandra.__version__)

# Connect to the VectorDB
def vector_db_conn():
    print("Connecting to the VectorDB.....")
   # This secure connect bundle is autogenerated when you donwload your SCB,
    # if yours is different update the file name below
    cloud_config= { 'secure_connect_bundle': os.getenv("SCB_path")
                }

    # This token json file is autogenerated when you donwload your token,
    # if yours is different update the file name below
    with open(os.getenv("DB_TOKEN") ) as f:
        secrets = json.load(f)

    CLIENT_ID = secrets["clientId"]
    CLIENT_SECRET = secrets["secret"]

    auth_provider = PlainTextAuthProvider(CLIENT_ID, CLIENT_SECRET)
    cluster = Cluster(cloud=cloud_config, auth_provider=auth_provider)
    session = cluster.connect()
    return session



session = vector_db_conn()
row = session.execute("select release_version from system.local").one()
if row:
    # print(row[0])
    print("Connection to the VectorDB Done.!")
else:
    print("An error occurred.")


# Define the Gradient's Model Adapter for LLAMA-2
gradient_llm = GradientBaseModelLLM(
    base_model_slug="llama2-7b-chat",
    max_tokens=400,
)   

# Configure Gradient embeddings
embed_model = GradientEmbedding(
    gradient_access_token = os.getenv("GRADIENT_ACCESS_TOKEN"),
    gradient_workspace_id = os.getenv("GRADIENT_WORKSPACE_ID"),
    gradient_model_slug="bge-large",
)

service_context = ServiceContext.from_defaults(
    llm = gradient_llm,
    embed_model = embed_model,
    chunk_size=256,
)
set_global_service_context(service_context)

# Load the PDFs
documents = SimpleDirectoryReader("Documents").load_data()
print(f"Loaded {len(documents)} document(s).")
print("________"*10)
print()
# Setup and Query Index
index = VectorStoreIndex.from_documents(documents,
                                        service_context=service_context)
query_engine = index.as_query_engine()

# Questions-Answer on Cassandra Paper
que1 = "What is Cassandra?"
response = query_engine.query(que1)
print("Question: ", que1)
print("Answer  : ", response)
print("________"*10)

# Questions-Answer on Attention Paper
que2 = "What is Positional Encoding?"
response2 = query_engine.query(que2)
print("Question: ", que2)
print("Answer  : ", response2)
print("________"*10)