import time
import gradio as gr

import asyncio
from pymongo import MongoClient
from langchain.vectorstores import MongoDBAtlasVectorSearch
from langchain.embeddings import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.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 = 'movies'
collection = client[db_name][collection_name]

## Create a collection and insert data
print ('Creating collection movies')
client[db_name].create_collection(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)

vector_store = MongoDBAtlasVectorSearch(embedding=OpenAIEmbeddings(), collection=collection, index_name='default', text_key='plot', embedding_key='plot_embedding')

def get_movies(message, history):
    movies = vector_store.similarity_search(message, 3)
    for movie in movies:
        for i in range(len(movie.metadata['title'])):
            time.sleep(0.05)
            yield "Movie " + i + " :  Title - " + movie.metadata['title'][: i+1]

demo = gr.ChatInterface(get_movies, examples=["What movies are scary?", "Find me a comedy", "Movies for kids"], title="Movies Atlas Vector Search", submit_btn="Search").queue()

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