import pymongo
import network_constants as netconst

class MongoRepository:
    # TODO extract to docker env var
    _username = 'root'
    _password = 'root'

    _collection: pymongo.collection.Collection = None
    _mongo_client: pymongo.MongoClient = None

    def __init__(self, username=_username, password=_password):
        self._mongo_client = pymongo.MongoClient(f"mongodb://{username}:{password}@{netconst.MONGO_DB_HOSTNAME}:{netconst.MONGO_DB_PORT}/")
        database = self._mongo_client['traceRetrievalDB']
        self._collection = database['traces']

    def insert_trace(self, content: dict):
        self._collection.insert_one(content)

    def get_traces(self, selection: dict = {}, projection: dict = {'_': 0}) -> pymongo.cursor.Cursor:
        return self._collection.find(selection, projection)

    def close_connection(self):
        self._mongo_client.close()
        self._collection = None
        self._mongo_client = None
