import os
import os.path
import pymongo
from pymongo.mongo_client import MongoClient
import gradio as gr
import certifi

MONGODB_ATLAS_DB_PASSWORD = os.environ['MONGODB_ATLAS_DB_PASSWORD']


from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch

from llama_index import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    load_index_from_storage,
)


mongo_uri = (
    f"mongodb+srv://arkiitkgp:{MONGODB_ATLAS_DB_PASSWORD}@genaicluster0.fgmvvsx.mongodb.net/?retryWrites=true&w=majority"
)
mongodb_client = pymongo.MongoClient(mongo_uri, tlsCAFile=certifi.where())

# Send a ping to confirm a successful connection
try:
    mongodb_client.admin.command('ping')
    print("Pinged your deployment. You successfully connected to MongoDB!")
except Exception as e:
    print(e)


store = MongoDBAtlasVectorSearch(mongodb_client, db_name='testdb1', collection_name='dummyIndex', index_name='vector_index')
# storage_context = StorageContext.from_defaults(vector_store=store)
# documents = SimpleDirectoryReader(
#         "temp", 
#         file_metadata=get_metadata_from_filename, 
#         required_exts=['.pdf', '.docx'],
#         recursive=True,
#     ).load_data()

# construct index
# index = VectorStoreIndex.from_documents(
#     documents, storage_context=storage_context
# )

index = VectorStoreIndex.from_vector_store(store)


def get_answer(query):
    response = index.as_query_engine(streaming=True).query(query)
    # response.print_response_stream()
    return response



###### GRADIO #################################

classes = ['Class 10', 'Class 9']
subjects = ['Science']

def get_streaming_answer(input_query):
    r = get_answer(input_query)
    ans = ""
    for new_tokens in r.response_gen:
        ans += new_tokens
        yield ans

with gr.Blocks() as demo:
    gr.Markdown("""# NCERT Tutor \n ### Type your question...""")
    with gr.Tab("For Students"):
        choose_class = gr.Dropdown(label= "Class", choices=classes)
        choose_subject = gr.Dropdown(label="Subject", choices=subjects)
        question_input = gr.Textbox(label="Enter question...")
        submit_button = gr.Button("Ask")
        response_output = gr.Textbox(label="Answer...", lines=5)
        
    with gr.Tab("For Teachers"):
        gr.Markdown("Coming soon...")

    submit_button.click(fn=get_streaming_answer, inputs=question_input, outputs=response_output)
  
demo.launch(share=True, debug=True)