import json
import socket
from typing import Any

import paho.mqtt.client as mqtt
from rich.console import Console

from .config import MqttConfig

console = Console()


class MqttClient:
    """
    A wrapper around paho.mqtt.client for handling MQTT connections and publishing.

    This class supports usage as a context manager for automatic connection and disconnection.
    """

    def __init__(self, config: MqttConfig) -> None:
        """
        Initialize the MQTT client.

        Args:
            config (MqttConfig): The MQTT configuration.
        """
        self.config = config
        self.client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)

        if self.config.username:
            self.client.username_pw_set(self.config.username, self.config.password)

    def connect(self) -> None:
        """
        Connect to the MQTT broker and start the network loop.

        Raises:
             socket.error: If the connection fails.
             OSError: If the connection fails.
        """
        try:
            self.client.connect(self.config.host, self.config.port, 60)
            self.client.loop_start()
        except (socket.error, OSError) as e:
            console.print(f"[bold red]Connection failed:[/bold red] {e}")
            raise

    def disconnect(self) -> None:
        """
        Stop the network loop and disconnect from the MQTT broker.
        """
        self.client.loop_stop()
        self.client.disconnect()

    def publish(self, topic: str, payload: Any, retain: bool = False) -> Any:
        """
        Publish a message to an MQTT topic.

        If the payload is not a string, it is serialized to JSON.
        This method blocks until the message is published.

        Args:
            topic (str): The MQTT topic to publish to.
            payload (Any): The message payload.
            retain (bool): Whether to retain the message. Defaults to False.

        Returns:
            mqtt.MQTTMessageInfo: Information about the published message.
        """
        if not isinstance(payload, str):
            payload = json.dumps(payload)

        info = self.client.publish(topic, payload, retain=retain)
        info.wait_for_publish()
        return info

    def __enter__(self) -> "MqttClient":
        self.connect()
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        self.disconnect()
