from dotenv import load_dotenv
import pandas as pd
import os
import pymongo
from langchain_mongodb.vectorstores import MongoDBAtlasVectorSearch
from langchain_huggingface import HuggingFaceEmbeddings

__author__ = "Chirag Kamble"

class DBConnect:
    """
    Class to connect to the database
    """
    @staticmethod
    def connect_db():
        """
        Static method to connect to the database and create a vector store
        :return: mongodb_vector_store: MongoDB Atlas Vector Store instance connected to the required mongodb collection
        :return: movies: dataframe containing all movies in the database
        """
        load_dotenv()

        mongodb_connection_url = os.getenv("MONGODB_CONNECTION_URL")
        mongodb_db_name: str = os.getenv("MONGODB_DB_NAME")
        mongodb_collection_name: str = os.getenv("MONGODB_COLLECTION_NAME")
        mongodb_vector_index: str = os.getenv("MONGODB_VECTOR_INDEX_NAME")
        text_key: str = os.getenv("TEXT_KEY")
        embedding_key: str = os.getenv("EMBEDDING_KEY")
        relevance_score_fn = os.getenv("RELEVANCE_SCORE_FN")

        client = pymongo.MongoClient(mongodb_connection_url)
        db = client[mongodb_db_name]
        collection = db[mongodb_collection_name]

        mongodb_vector_store = MongoDBAtlasVectorSearch(collection=collection,
                                                        embedding=HuggingFaceEmbeddings(),
                                                        index_name=mongodb_vector_index,
                                                        relevance_score_fn=relevance_score_fn,
                                                        text_key=text_key,
                                                        embedding_key=embedding_key,
                                                        )

        movies_docs = collection.find()
        movies = pd.DataFrame(movies_docs)

        return mongodb_vector_store, movies