import paho.mqtt.client as mqtt
import json
import logging
from typing import Callable

class MQTTClient:
    def __init__(self, on_message_callback: Callable[[str, dict], None]):
        self.on_message_callback = on_message_callback
        self.client = mqtt.Client()
        self.client.on_connect = self._on_connect
        self.client.on_disconnect = self._on_disconnect
        self.client.on_message = self._on_message
        self.connected = False
        self._pending_topics: list = []

    def connect(self, broker: str, port: int = 1883):
        self.disconnect()
        try:
            logging.info(f"MQTT正在连接 {broker}:{port}")
            self.client.connect(broker, port, 60)
            self.client.loop_start()
        except Exception as e:
            logging.error(f"MQTT连接失败 {broker}:{port}: {e}")

    def disconnect(self):
        self.connected = False
        try:
            self.client.loop_stop()
            self.client.disconnect()
        except Exception:
            pass

    def subscribe(self, topics):
        self._pending_topics.extend([t for t in topics if t not in self._pending_topics])
        if not self.connected:
            return
        for topic in topics:
            self.client.subscribe(topic)

    def unsubscribe(self, topic):
        self._pending_topics = [t for t in self._pending_topics if t != topic]
        if not self.connected:
            return
        self.client.unsubscribe(topic)

    def _on_connect(self, client, userdata, flags, rc):
        if rc == 0:
            self.connected = True
            logging.info(f"MQTT已连接，正在订阅 {len(self._pending_topics)} 个topic")
            for topic in self._pending_topics:
                client.subscribe(topic)
        else:
            logging.warning(f"MQTT broker拒绝连接，返回码: {rc}")

    def _on_disconnect(self, client, userdata, rc):
        self.connected = False
        if rc != 0:
            logging.warning(f"MQTT意外断开，rc={rc}")

    def _on_message(self, client, userdata, msg):
        try:
            payload_str = msg.payload.decode('utf-8')
            if not payload_str.strip():
                return
            payload = json.loads(payload_str)
            if payload is not None:
                self.on_message_callback(msg.topic, payload)
            else:
                logging.debug("收到 null payload，忽略")
        except Exception as e:
            logging.exception(f"MQTT消息解析错误")