import paho.mqtt.client as paho
from edge_engine.streamio.datastream.datastreamwrapper import DataStreamWrapper
from edge_engine.common.logsetup import logger
from uuid import uuid4
import traceback
class MQTT(DataStreamWrapper):
    @staticmethod
    def on_connect(client, userdata, flags, rc):
        logger.info("Connection returned with result code:" + str(rc))

    @staticmethod
    def on_disconnect(client, userdata, rc):
        logger.info("Disconnection returned result:" + str(rc))

    @staticmethod
    def on_subscribe(client, userdata, mid, granted_qos):
        logger.debug("Subscribing MQTT {} {} {} {}".format(client, userdata, mid, granted_qos))

    def on_message(self, client, userdata, msg):
        logger.debug("Received message, topic:" + msg.topic + "payload:" + str(msg.payload))
        if self.subscribe_hook is not None:
            self.subscribe_hook(msg.payload.decode())

    def __init__(self, broker, port, topic, qos=2, subscribe_hook=None, publish_hook=None):
        super().__init__()
        self.broker = broker
        self.port = int(port)
        self.topic = topic
        self.client_name = "{}".format(uuid4())
        self.client = paho.Client(self.client_name)
        self.client.on_connect = self.on_connect
        self.client.on_disconnect = self.on_disconnect
        self.client.on_subscribe = self.on_subscribe
        self.client.on_message = self.on_message
        self.client.connect(host=self.broker, port=self.port)
        self.subscribe_hook = subscribe_hook
        self.publish_hook = publish_hook
        self.qos = qos

    def subscribe(self, hook=None):
        if hook is not None:
            self.subscribe_hook =hook
        self.client.subscribe((self.topic, self.qos))
        self.client.loop_forever()

    def publish(self, data):
        try:
            if self.publish_hook is not None:
                data = self.publish_hook(data)
            self.client.publish(self.topic, data)
        except Exception as e:
            logger.error(e)
            logger.error(traceback.format_exc())
