import time
import gradio as gr
import os
import asyncio
from pymongo import MongoClient
from langchain_community.vectorstores import MongoDBAtlasVectorSearch
from langchain_openai import OpenAIEmbeddings
from langchain_community.llms import OpenAI
# from langchain_community.prompts import PromptTemplate
# from langchain.chains import LLMChain
import json


## Connect to MongoDB Atlas local cluster
MONGODB_ATLAS_CLUSTER_URI = os.getenv('MONGODB_ATLAS_CLUSTER_URI')
client = MongoClient(MONGODB_ATLAS_CLUSTER_URI)
db_name = 'sample_mflix'
collection_name = 'embedded_movies'
collection = client[db_name][collection_name]

## Create a vector search index
print ('Creating vector search index')
# collection.create_search_index(model={"definition": {"mappings":{
#   "dynamic":True,
#   "fields": {
#     "plot_embedding": {
#       "type": "knnVector",
#       "dimensions": 1536,
#       "similarity": "euclidean"
#     }
#   }
# }}, "name":'default'})

# sleep for minute
# print ('Waiting for vector index on field "embedding" to be created')
# time.sleep(60)

try:
    vector_store = MongoDBAtlasVectorSearch(embedding=OpenAIEmbeddings(), collection=collection, index_name='vector_index', text_key='plot', embedding_key='plot_embedding')

except:
    # If open ai key is wrong
    print ('Open AI key is wrong')
    vector_store = None

def get_movies(message, history):
    # Use AsyncIO to run the similarity search in the background
   # movies =  vector_store.similarity_search(message, 3)
    print ('Searching for: ' + message)
    try:
        movies =  vector_store.similarity_search(message, 3)
        retrun_text = ''
        for movie in movies:
            retrun_text = retrun_text + 'Title : ' +  movie.metadata['title'] + '\n------------\n' + 'Plot: ' + movie.page_content + '\n\n'

        for i in range(len(retrun_text)):
            time.sleep(0.05)
            yield "Found: " + "\n\n" + retrun_text[: i+1]
    except:
        yield "Please clone the repo and add your open ai key as well as your MongoDB Atlas UR in the Secret Section of you Space\n OPENAI_API_KEY (your Open AI key) and MONGODB_ATLAS_CLUSTER_URI (0.0.0.0/0 whitelisted instance with Vector index created) \n\n For more information : https://mongodb.com/products/platform/atlas-vector-search"
    

demo = gr.ChatInterface(get_movies, examples=["What movies are scary?", "Find me a comedy", "Movies for kids"],description="Uses the sample_mflix.embedded_movies collection from [sample_data](https://www.mongodb.com/docs/atlas/sample-data/sample-mflix/#sample_mflix.embedded_movies) require a [vector index](https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-tutorial/#create-the-atlas-vector-search-index) on `plot_embedding` field", title="Movies Atlas Vector Search", submit_btn="Search").queue()

if __name__ == "__main__":
    demo.launch()