from edge_engine.streamio.datastream.datastreamwrapper import DataStreamWrapper
from pymongo import MongoClient


class MongoDataStreamOut(DataStreamWrapper):

    def __init__(self, host, port, dbname, collection, keys, authsource,username=None,password=None, publish_hook=None):
        super().__init__()
        self.host = host
        self.port = port
        self.dbname = dbname
        self.username = username
        self.password = password
        self.collection = collection
        self.publish_hook = publish_hook
        self.mongo = MongoClient(host=host,
                                 port=int(port),username=self.username,password=self.password)
        self.db = self.mongo[dbname]
        self.keys = keys
        self.authsource = authsource

    def subscribe(self, hook=None):
        pass

    def publish(self, data):
        if self.publish_hook is not None:
            data = self.publish_hook(data)
        fin_dat = {}
        for k, v in data.items():
            if k in self.keys:
                fin_dat[k] = v
        self.db[self.collection].insert(fin_dat)
