import json
import logging
import random
import time
from datetime import datetime

from paho.mqtt import client as mqtt_client


class MQTTClient:
    def __init__(self, broker='218.3.126.49', port=1883, topic="wybb/mqtt",
                 username='emqx', password='public',client_id=''):
        self.BROKER = broker
        self.PORT = port
        self.TOPIC = topic
        # generate client ID with pub prefix randomly
        self.CLIENT_ID = client_id if client_id!='' else f'wybb_{datetime.now().strftime("%Y%m%d_%H%M%S")}'

        self.USERNAME = username
        self.PASSWORD = password

        self.FIRST_RECONNECT_DELAY = 1
        self.RECONNECT_RATE = 2
        self.MAX_RECONNECT_COUNT = 10
        self.MAX_RECONNECT_DELAY = 60

        self.FLAG_EXIT = False
        self.client = None

    def on_connect(self, client, userdata, flags, rc, properties=None):
        if rc == 0 and client.is_connected():
            print("Connected to MQTT Broker!")
        else:
            print(f'Failed to connect, return code {rc}')

    def on_disconnect(self, client, userdata, flags, rc, properties=None):
        logging.info("Disconnected with result code: %s", rc)
        reconnect_count, reconnect_delay = 0, self.FIRST_RECONNECT_DELAY
        while reconnect_count < self.MAX_RECONNECT_COUNT:
            logging.info("Reconnecting in %d seconds...", reconnect_delay)
            time.sleep(reconnect_delay)

            try:
                client.reconnect()
                logging.info("Reconnected successfully!")
                return
            except Exception as err:
                logging.error("%s. Reconnect failed. Retrying...", err)

            reconnect_delay *= self.RECONNECT_RATE
            reconnect_delay = min(reconnect_delay, self.MAX_RECONNECT_DELAY)
            reconnect_count += 1
        logging.info("Reconnect failed after %s attempts. Exiting...", reconnect_count)
        self.FLAG_EXIT = True

    def connect(self):
        self.client = mqtt_client.Client(mqtt_client.CallbackAPIVersion.VERSION2, self.CLIENT_ID)
        self.client.username_pw_set(self.USERNAME, self.PASSWORD)
        self.client.on_connect = self.on_connect
        self.client.on_disconnect = self.on_disconnect
        self.client.connect(self.BROKER, self.PORT, keepalive=120)
        return self.client

    def publish(self, msg:str):
        if not self.client:
            logging.error("MQTT client is not initialized!")
            return

        if not self.client.is_connected():
            logging.error("publish: MQTT client is not connected!")
            time.sleep(1)
            return
        result = self.client.publish(self.TOPIC, msg)
        # result: [0, 1]
        status = result[0]
        if status == 0:
            # logging.info(f'Send topic [{self.TOPIC}] to `{msg}`  ')
            pass
        else:
            logging.info(f'Failed to send topic [{self.TOPIC}]')

    def start(self):
        self.client = self.connect()
        self.client.loop_start()
        time.sleep(1)
        if self.client.is_connected():
            logging.info("MQTT client is connected!")
        else:
            self.client.loop_stop()
            return False
        return True

    def stop(self):
        if self.client:
            self.client.loop_stop()
            self.client.disconnect()


# 示例用法
if __name__ == '__main__':
    mq = MQTTClient()
    if mq.start():
        msg = {"a": 123}
        mq.publish(msg)
        time.sleep(2)  # 等待消息发送完成
        mq.stop()
    print("==========")