"""MQTT client for Heiman device communication."""

import asyncio
import json
import logging
from typing import Any, Callable, Optional

from aiomqtt import Client, MqttError, Message

from .exceptions import HeimanMQTTError
from .models import HeimanDevice

_LOGGER = logging.getLogger(__name__)


class HeimanMqttClient:
    """MQTT client for real-time device communication.
    
    This client connects to Heiman MQTT broker and subscribes to device topics
    for receiving real-time updates.
    """

    def __init__(
        self,
        mqtt_broker: str,
        mqtt_port: int = 1883,
        mqtt_username: Optional[str] = None,
        mqtt_password: Optional[str] = None,
        client_id: Optional[str] = None,
        use_tls: bool = False,
    ) -> None:
        """Initialize MQTT client.
        
        Args:
            mqtt_broker: MQTT broker hostname
            mqtt_port: MQTT broker port
            mqtt_username: MQTT username (usually user ID)
            mqtt_password: MQTT password (usually access token)
            client_id: MQTT client ID
            use_tls: Whether to use TLS encryption
        """
        self.mqtt_broker = mqtt_broker
        self.mqtt_port = mqtt_port
        self.mqtt_username = mqtt_username
        self.mqtt_password = mqtt_password
        self.client_id = client_id or f"heiman_ha_{id(self)}"
        self.use_tls = use_tls
        
        self._client: Optional[Client] = None
        self._running = False
        self._subscriptions: set[str] = set()
        self._message_callbacks: dict[str, list[Callable[[str, Any], None]]] = {}
        self._connect_task: Optional[asyncio.Task] = None
        self._disconnect_event = asyncio.Event()

    async def connect(self) -> None:
        """Connect to MQTT broker.
        
        Raises:
            HeimanMQTTError: If connection fails
        """
        if self._running:
            _LOGGER.debug("MQTT client already running")
            return
        
        try:
            self._client = Client(
                hostname=self.mqtt_broker,
                port=self.mqtt_port,
                username=self.mqtt_username,
                password=self.mqtt_password,
                client_id=self.client_id,
                tls_context={"ca_certs": None} if self.use_tls else None,
            )
            
            self._running = True
            self._disconnect_event.clear()
            
            # Start message processing loop
            asyncio.create_task(self._message_loop())
            
            _LOGGER.info(f"Connected to MQTT broker {self.mqtt_broker}")
            
        except MqttError as err:
            raise HeimanMQTTError(
                message=f"Failed to connect to MQTT broker: {str(err)}",
            ) from err
        except Exception as err:
            raise HeimanMQTTError(
                message=f"Unexpected error connecting to MQTT: {str(err)}",
            ) from err

    async def disconnect(self) -> None:
        """Disconnect from MQTT broker."""
        if not self._running:
            return
        
        self._running = False
        self._disconnect_event.set()
        
        if self._client and not self._client._client.is_disconnected():
            try:
                await self._client.disconnect()
            except Exception as err:
                _LOGGER.warning(f"Error during MQTT disconnect: {err}")
        
        _LOGGER.info("Disconnected from MQTT broker")

    async def _message_loop(self) -> None:
        """Process incoming MQTT messages."""
        if not self._client:
            return
        
        try:
            async with self._client:
                # Subscribe to all registered topics
                for topic in self._subscriptions:
                    await self._client.subscribe(topic)
                    _LOGGER.debug(f"Subscribed to topic: {topic}")
                
                # Process messages
                async for message in self._client.messages:
                    if not self._running:
                        break
                    
                    await self._handle_message(message)
                    
        except MqttError as err:
            if self._running:
                _LOGGER.error(f"MQTT error in message loop: {err}")
        except Exception as err:
            if self._running:
                _LOGGER.error(f"Unexpected error in message loop: {err}")

    async def _handle_message(self, message: Message) -> None:
        """Handle incoming MQTT message.
        
        Args:
            message: MQTT message object
        """
        topic = message.topic.value
        payload = message.payload.decode("utf-8", errors="ignore")
        
        _LOGGER.debug(f"Received MQTT message on topic {topic}: {payload[:200]}")
        
        try:
            data = json.loads(payload) if payload else {}
        except json.JSONDecodeError:
            _LOGGER.warning(f"Failed to parse MQTT payload: {payload[:100]}")
            data = {"raw": payload}
        
        # Call registered callbacks
        callbacks = self._message_callbacks.get(topic, [])
        for callback in callbacks:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(topic, data)
                else:
                    callback(topic, data)
            except Exception as err:
                _LOGGER.error(f"Error in MQTT callback for topic {topic}: {err}")

    def subscribe(
        self,
        topic: str,
        callback: Callable[[str, Any], None],
    ) -> None:
        """Subscribe to MQTT topic with callback.
        
        Args:
            topic: Topic pattern to subscribe to
            callback: Callback function to invoke on message
        """
        self._subscriptions.add(topic)
        
        if topic not in self._message_callbacks:
            self._message_callbacks[topic] = []
        
        self._message_callbacks[topic].append(callback)
        _LOGGER.debug(f"Registered callback for topic: {topic}")

    def unsubscribe(self, topic: str) -> None:
        """Unsubscribe from MQTT topic.
        
        Args:
            topic: Topic to unsubscribe
        """
        if topic in self._subscriptions:
            self._subscriptions.discard(topic)
        
        if topic in self._message_callbacks:
            del self._message_callbacks[topic]
        
        _LOGGER.debug(f"Unsubscribed from topic: {topic}")

    async def publish(
        self,
        topic: str,
        payload: dict[str, Any],
        qos: int = 1,
        retain: bool = False,
    ) -> None:
        """Publish message to MQTT topic.
        
        Args:
            topic: Topic to publish to
            payload: Message payload (will be JSON encoded)
            qos: Quality of Service level
            retain: Whether to retain the message
            
        Raises:
            HeimanMQTTError: If publish fails
        """
        if not self._running or not self._client:
            raise HeimanMQTTError(message="Not connected to MQTT broker")
        
        try:
            payload_str = json.dumps(payload)
            await self._client.publish(topic, payload=payload_str, qos=qos, retain=retain)
            _LOGGER.debug(f"Published to {topic}: {payload_str[:200]}")
        except MqttError as err:
            raise HeimanMQTTError(
                message=f"Failed to publish message: {str(err)}",
            ) from err

    def get_topic_for_device(self, device_id: str, message_type: str = "report") -> str:
        """Get MQTT topic for a specific device.
        
        Args:
            device_id: Device ID
            message_type: Type of message (report, control, response)
            
        Returns:
            MQTT topic string
        """
        # Topic format based on Heiman protocol
        return f"/device/{device_id}/{message_type}"

    @property
    def is_connected(self) -> bool:
        """Check if MQTT client is connected."""
        return self._running and self._client is not None
