import paho.mqtt.client as mqtt
import logging
import time
from threading import Thread, Event

logger = logging.getLogger(__name__)

class ConnectionManager:
    def __init__(self):
        self.client = None
        self.connected = False
        self._stop_event = Event()
        self._heartbeat_thread = None
        self._message_handler = None

    def set_message_handler(self, handler):
        self._message_handler = handler

    def connect(self, broker_ip: str, client_id: str, port: int = 1883, username: str = None, password: str = None):
        """
        Connects to the MQTT Broker.
        """
        if self.client:
            self.disconnect()

        # Create a new MQTT Client instance
        # Protocol V5 is recommended for modern deployments, but V3.1.1 is standard. 
        # Using default (usually 3.1.1 or 5.0 depending on paho version/negotiation).
        self.client = mqtt.Client(client_id=client_id, protocol=mqtt.MQTTv311)
        
        if username and password:
            self.client.username_pw_set(username, password)

        # Set callbacks
        self.client.on_connect = self._on_connect
        self.client.on_disconnect = self._on_disconnect
        self.client.on_message = self._on_message
        self.client.on_publish = self._on_publish

        logger.info(f"Connecting to {broker_ip}:{port} as {client_id}...")
        try:
            self.client.connect(broker_ip, port, keepalive=60)
            self.client.loop_start() # Start the network loop in a background thread
        except Exception as e:
            logger.error(f"Failed to connect: {e}")
            raise

    def disconnect(self):
        """
        Disconnects from the broker.
        """
        self._stop_heartbeat()
        if self.client:
            self.client.loop_stop()
            self.client.disconnect()
            self.client = None
            self.connected = False
            logger.info("Disconnected.")

    def _on_connect(self, client, userdata, flags, rc):
        if rc == 0:
            self.connected = True
            logger.info("Connected successfully.")
            self._start_heartbeat()
            # Subscribe to services if needed, but usually caller handles subscription 
            # or we do it here if we know the topic.
            # For now, we rely on caller or hardcoded.
            # Let's subscribe to everything for this simulator or specific service topic
            # client.subscribe("#") 
            # Ideally, we should let the ServiceHandler subscribe.
        else:
            self.connected = False
            logger.error(f"Connection failed with result code {rc}")

    def _on_disconnect(self, client, userdata, rc):
        self.connected = False
        self._stop_heartbeat()
        if rc != 0:
            logger.warning("Unexpected disconnection.")
        else:
            logger.info("Disconnected gracefully.")

    def _on_message(self, client, userdata, msg):
        logger.debug(f"Received message on {msg.topic}: {msg.payload}")
        if self._message_handler:
            self._message_handler(msg.topic, msg.payload)

    def _on_publish(self, client, userdata, mid):
        pass

    def _start_heartbeat(self):
        self._stop_event.clear()
        self._heartbeat_thread = Thread(target=self._heartbeat_loop, daemon=True)
        self._heartbeat_thread.start()

    def _stop_heartbeat(self):
        self._stop_event.set()
        if self._heartbeat_thread:
            self._heartbeat_thread.join(timeout=1.0)
            self._heartbeat_thread = None

    def _heartbeat_loop(self):
        logger.info("Heartbeat loop started.")
        while not self._stop_event.is_set():
            if self.client and self.connected:
                # Actual topic/payload depends on DJI API specs. 
                # For now, we simulate a ping or minimal publish.
                # Usually MQTT keepalive handles PINGREQ/PINGRESP automatically.
                # But sometimes application-level heartbeat is needed.
                # We will rely on Paho's built-in keepalive for network ping, 
                # but if an app-level heartbeat is needed (e.g. reporting state), we do it here.
                # For now, let's just log or sleep. 
                pass 
            time.sleep(5)

    def is_connected(self):
        return self.connected