"""
Support for MQTT message handling.

For more details about this component, please refer to the documentation at
https://home-assistant.io/components/mqtt/
"""
import asyncio
from itertools import groupby
import json
import logging
from operator import attrgetter
import os
import socket
import ssl
import time
from typing import Any, Callable, List, Optional, Union, cast  # noqa: F401

import attr
import requests.certs
import voluptuous as vol

from homeassistant import config_entries
from homeassistant.const import (
    CONF_DEVICE, CONF_NAME, CONF_PASSWORD, CONF_PAYLOAD, CONF_PORT,
    CONF_PROTOCOL, CONF_USERNAME, CONF_VALUE_TEMPLATE,
    EVENT_HOMEASSISTANT_STOP)
from homeassistant.core import Event, ServiceCall, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import config_validation as cv, template
from homeassistant.helpers.entity import Entity
from homeassistant.helpers.typing import (
    ConfigType, HomeAssistantType, ServiceDataType)
from homeassistant.loader import bind_hass
from homeassistant.setup import async_prepare_setup_platform
from homeassistant.util.async_ import (
    run_callback_threadsafe, run_coroutine_threadsafe)
from homeassistant.util.logging import catch_log_exception

# Loading the config flow file will register the flow
from . import config_flow  # noqa pylint: disable=unused-import
from .const import CONF_BROKER, CONF_DISCOVERY, DEFAULT_DISCOVERY
from .server import HBMQTT_CONFIG_SCHEMA

REQUIREMENTS = ['paho-mqtt==1.4.0']

_LOGGER = logging.getLogger(__name__)

DOMAIN = 'mqtt'

DATA_MQTT = 'mqtt'
DATA_MQTT_CONFIG = 'mqtt_config'
DATA_MQTT_HASS_CONFIG = 'mqtt_hass_config'

SERVICE_PUBLISH = 'publish'

CONF_EMBEDDED = 'embedded'

CONF_CLIENT_ID = 'client_id'
CONF_DISCOVERY_PREFIX = 'discovery_prefix'
CONF_KEEPALIVE = 'keepalive'
CONF_CERTIFICATE = 'certificate'
CONF_CLIENT_KEY = 'client_key'
CONF_CLIENT_CERT = 'client_cert'
CONF_TLS_INSECURE = 'tls_insecure'
CONF_TLS_VERSION = 'tls_version'

CONF_BIRTH_MESSAGE = 'birth_message'
CONF_WILL_MESSAGE = 'will_message'

CONF_STATE_TOPIC = 'state_topic'
CONF_COMMAND_TOPIC = 'command_topic'
CONF_AVAILABILITY_TOPIC = 'availability_topic'
CONF_PAYLOAD_AVAILABLE = 'payload_available'
CONF_PAYLOAD_NOT_AVAILABLE = 'payload_not_available'
CONF_JSON_ATTRS_TOPIC = 'json_attributes_topic'
CONF_QOS = 'qos'
CONF_RETAIN = 'retain'

CONF_UNIQUE_ID = 'unique_id'
CONF_IDENTIFIERS = 'identifiers'
CONF_CONNECTIONS = 'connections'
CONF_MANUFACTURER = 'manufacturer'
CONF_MODEL = 'model'
CONF_SW_VERSION = 'sw_version'
CONF_VIA_HUB = 'via_hub'

PROTOCOL_31 = '3.1'
PROTOCOL_311 = '3.1.1'

DEFAULT_PORT = 1883
DEFAULT_KEEPALIVE = 60
DEFAULT_QOS = 0
DEFAULT_RETAIN = False
DEFAULT_PROTOCOL = PROTOCOL_311
DEFAULT_DISCOVERY_PREFIX = 'homeassistant'
DEFAULT_TLS_PROTOCOL = 'auto'
DEFAULT_PAYLOAD_AVAILABLE = 'online'
DEFAULT_PAYLOAD_NOT_AVAILABLE = 'offline'

ATTR_TOPIC = 'topic'
ATTR_PAYLOAD = 'payload'
ATTR_PAYLOAD_TEMPLATE = 'payload_template'
ATTR_QOS = CONF_QOS
ATTR_RETAIN = CONF_RETAIN
ATTR_DISCOVERY_HASH = 'discovery_hash'

MAX_RECONNECT_WAIT = 300  # seconds


def valid_topic(value: Any) -> str:
    """Validate that this is a valid topic name/filter."""
    value = cv.string(value)
    try:
        raw_value = value.encode('utf-8')
    except UnicodeError:
        raise vol.Invalid("MQTT topic name/filter must be valid UTF-8 string.")
    if not raw_value:
        raise vol.Invalid("MQTT topic name/filter must not be empty.")
    if len(raw_value) > 65535:
        raise vol.Invalid("MQTT topic name/filter must not be longer than "
                          "65535 encoded bytes.")
    if '\0' in value:
        raise vol.Invalid("MQTT topic name/filter must not contain null "
                          "character.")
    return value


def valid_subscribe_topic(value: Any) -> str:
    """Validate that we can subscribe using this MQTT topic."""
    value = valid_topic(value)
    for i in (i for i, c in enumerate(value) if c == '+'):
        if (i > 0 and value[i - 1] != '/') or \
                (i < len(value) - 1 and value[i + 1] != '/'):
            raise vol.Invalid("Single-level wildcard must occupy an entire "
                              "level of the filter")

    index = value.find('#')
    if index != -1:
        if index != len(value) - 1:
            # If there are multiple wildcards, this will also trigger
            raise vol.Invalid("Multi-level wildcard must be the last "
                              "character in the topic filter.")
        if len(value) > 1 and value[index - 1] != '/':
            raise vol.Invalid("Multi-level wildcard must be after a topic "
                              "level separator.")

    return value


def valid_publish_topic(value: Any) -> str:
    """Validate that we can publish using this MQTT topic."""
    value = valid_topic(value)
    if '+' in value or '#' in value:
        raise vol.Invalid("Wildcards can not be used in topic names")
    return value


def validate_device_has_at_least_one_identifier(value: ConfigType) -> \
        ConfigType:
    """Validate that a device info entry has at least one identifying value."""
    if not value.get(CONF_IDENTIFIERS) and not value.get(CONF_CONNECTIONS):
        raise vol.Invalid("Device must have at least one identifying value in "
                          "'identifiers' and/or 'connections'")
    return value


_VALID_QOS_SCHEMA = vol.All(vol.Coerce(int), vol.In([0, 1, 2]))

CLIENT_KEY_AUTH_MSG = 'client_key and client_cert must both be present in ' \
                      'the MQTT broker configuration'

MQTT_WILL_BIRTH_SCHEMA = vol.Schema({
    vol.Required(ATTR_TOPIC): valid_publish_topic,
    vol.Required(ATTR_PAYLOAD, CONF_PAYLOAD): cv.string,
    vol.Optional(ATTR_QOS, default=DEFAULT_QOS): _VALID_QOS_SCHEMA,
    vol.Optional(ATTR_RETAIN, default=DEFAULT_RETAIN): cv.boolean,
}, required=True)

CONFIG_SCHEMA = vol.Schema({
    DOMAIN: vol.Schema({
        vol.Optional(CONF_CLIENT_ID): cv.string,
        vol.Optional(CONF_KEEPALIVE, default=DEFAULT_KEEPALIVE):
            vol.All(vol.Coerce(int), vol.Range(min=15)),
        vol.Optional(CONF_BROKER): cv.string,
        vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
        vol.Optional(CONF_USERNAME): cv.string,
        vol.Optional(CONF_PASSWORD): cv.string,
        vol.Optional(CONF_CERTIFICATE): vol.Any('auto', cv.isfile),
        vol.Inclusive(CONF_CLIENT_KEY, 'client_key_auth',
                      msg=CLIENT_KEY_AUTH_MSG): cv.isfile,
        vol.Inclusive(CONF_CLIENT_CERT, 'client_key_auth',
                      msg=CLIENT_KEY_AUTH_MSG): cv.isfile,
        vol.Optional(CONF_TLS_INSECURE): cv.boolean,
        vol.Optional(CONF_TLS_VERSION, default=DEFAULT_TLS_PROTOCOL):
            vol.Any('auto', '1.0', '1.1', '1.2'),
        vol.Optional(CONF_PROTOCOL, default=DEFAULT_PROTOCOL):
            vol.All(cv.string, vol.In([PROTOCOL_31, PROTOCOL_311])),
        vol.Optional(CONF_EMBEDDED): HBMQTT_CONFIG_SCHEMA,
        vol.Optional(CONF_WILL_MESSAGE): MQTT_WILL_BIRTH_SCHEMA,
        vol.Optional(CONF_BIRTH_MESSAGE): MQTT_WILL_BIRTH_SCHEMA,
        vol.Optional(CONF_DISCOVERY, default=DEFAULT_DISCOVERY): cv.boolean,
        # discovery_prefix must be a valid publish topic because if no
        # state topic is specified, it will be created with the given prefix.
        vol.Optional(CONF_DISCOVERY_PREFIX,
                     default=DEFAULT_DISCOVERY_PREFIX): valid_publish_topic,
    }),
}, extra=vol.ALLOW_EXTRA)

SCHEMA_BASE = {
    vol.Optional(CONF_QOS, default=DEFAULT_QOS): _VALID_QOS_SCHEMA,
}

MQTT_AVAILABILITY_SCHEMA = vol.Schema({
    vol.Optional(CONF_AVAILABILITY_TOPIC): valid_subscribe_topic,
    vol.Optional(CONF_PAYLOAD_AVAILABLE,
                 default=DEFAULT_PAYLOAD_AVAILABLE): cv.string,
    vol.Optional(CONF_PAYLOAD_NOT_AVAILABLE,
                 default=DEFAULT_PAYLOAD_NOT_AVAILABLE): cv.string,
})

MQTT_ENTITY_DEVICE_INFO_SCHEMA = vol.All(vol.Schema({
    vol.Optional(CONF_IDENTIFIERS, default=list):
        vol.All(cv.ensure_list, [cv.string]),
    vol.Optional(CONF_CONNECTIONS, default=list):
        vol.All(cv.ensure_list, [vol.All(vol.Length(2), [cv.string])]),
    vol.Optional(CONF_MANUFACTURER): cv.string,
    vol.Optional(CONF_MODEL): cv.string,
    vol.Optional(CONF_NAME): cv.string,
    vol.Optional(CONF_SW_VERSION): cv.string,
    vol.Optional(CONF_VIA_HUB): cv.string,
}), validate_device_has_at_least_one_identifier)

MQTT_JSON_ATTRS_SCHEMA = vol.Schema({
    vol.Optional(CONF_JSON_ATTRS_TOPIC): valid_subscribe_topic,
})

MQTT_BASE_PLATFORM_SCHEMA = cv.PLATFORM_SCHEMA_2.extend(SCHEMA_BASE)

# Sensor type platforms subscribe to MQTT events
MQTT_RO_PLATFORM_SCHEMA = MQTT_BASE_PLATFORM_SCHEMA.extend({
    vol.Required(CONF_STATE_TOPIC): valid_subscribe_topic,
    vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
})

# Switch type platforms publish to MQTT and may subscribe
MQTT_RW_PLATFORM_SCHEMA = MQTT_BASE_PLATFORM_SCHEMA.extend({
    vol.Required(CONF_COMMAND_TOPIC): valid_publish_topic,
    vol.Optional(CONF_RETAIN, default=DEFAULT_RETAIN): cv.boolean,
    vol.Optional(CONF_STATE_TOPIC): valid_subscribe_topic,
    vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
})

# Service call validation schema
MQTT_PUBLISH_SCHEMA = vol.Schema({
    vol.Required(ATTR_TOPIC): valid_publish_topic,
    vol.Exclusive(ATTR_PAYLOAD, CONF_PAYLOAD): object,
    vol.Exclusive(ATTR_PAYLOAD_TEMPLATE, CONF_PAYLOAD): cv.string,
    vol.Optional(ATTR_QOS, default=DEFAULT_QOS): _VALID_QOS_SCHEMA,
    vol.Optional(ATTR_RETAIN, default=DEFAULT_RETAIN): cv.boolean,
}, required=True)


# pylint: disable=invalid-name
PublishPayloadType = Union[str, bytes, int, float, None]
SubscribePayloadType = Union[str, bytes]  # Only bytes if encoding is None
MessageCallbackType = Callable[[str, SubscribePayloadType, int], None]


def _build_publish_data(topic: Any, qos: int, retain: bool) -> ServiceDataType:
    """Build the arguments for the publish service without the payload."""
    data = {ATTR_TOPIC: topic}
    if qos is not None:
        data[ATTR_QOS] = qos
    if retain is not None:
        data[ATTR_RETAIN] = retain
    return data


@bind_hass
def publish(hass: HomeAssistantType, topic, payload, qos=None,
            retain=None) -> None:
    """Publish message to an MQTT topic."""
    hass.add_job(async_publish, hass, topic, payload, qos, retain)


@callback
@bind_hass
def async_publish(hass: HomeAssistantType, topic: Any, payload, qos=None,
                  retain=None) -> None:
    """Publish message to an MQTT topic."""
    data = _build_publish_data(topic, qos, retain)
    data[ATTR_PAYLOAD] = payload
    hass.async_create_task(
        hass.services.async_call(DOMAIN, SERVICE_PUBLISH, data))


@bind_hass
def publish_template(hass: HomeAssistantType, topic, payload_template,
                     qos=None, retain=None) -> None:
    """Publish message to an MQTT topic using a template payload."""
    data = _build_publish_data(topic, qos, retain)
    data[ATTR_PAYLOAD_TEMPLATE] = payload_template
    hass.services.call(DOMAIN, SERVICE_PUBLISH, data)


@bind_hass
async def async_subscribe(hass: HomeAssistantType, topic: str,
                          msg_callback: MessageCallbackType,
                          qos: int = DEFAULT_QOS,
                          encoding: str = 'utf-8'):
    """Subscribe to an MQTT topic.

    Call the return value to unsubscribe.
    """
    async_remove = await hass.data[DATA_MQTT].async_subscribe(
        topic, catch_log_exception(
            msg_callback, lambda topic, msg, qos:
            "Exception in {} when handling msg on '{}': '{}'".format(
                msg_callback.__name__, topic, msg)),
        qos, encoding)
    return async_remove


@bind_hass
def subscribe(hass: HomeAssistantType, topic: str,
              msg_callback: MessageCallbackType, qos: int = DEFAULT_QOS,
              encoding: str = 'utf-8') -> Callable[[], None]:
    """Subscribe to an MQTT topic."""
    async_remove = run_coroutine_threadsafe(
        async_subscribe(hass, topic, msg_callback, qos, encoding), hass.loop
    ).result()

    def remove():
        """Remove listener convert."""
        run_callback_threadsafe(hass.loop, async_remove).result()

    return remove


async def _async_setup_server(hass: HomeAssistantType, config: ConfigType):
    """Try to start embedded MQTT broker.

    This method is a coroutine.
    """
    conf = config.get(DOMAIN, {})  # type: ConfigType

    server = await async_prepare_setup_platform(
        hass, config, DOMAIN, 'server')

    if server is None:
        _LOGGER.error("Unable to load embedded server")
        return None

    success, broker_config = \
        await server.async_start(
            hass, conf.get(CONF_PASSWORD), conf.get(CONF_EMBEDDED))

    if not success:
        return None

    return broker_config


async def _async_setup_discovery(hass: HomeAssistantType, conf: ConfigType,
                                 hass_config: ConfigType,
                                 config_entry) -> bool:
    """Try to start the discovery of MQTT devices.

    This method is a coroutine.
    """
    discovery = await async_prepare_setup_platform(
        hass, hass_config, DOMAIN, 'discovery')

    if discovery is None:
        _LOGGER.error("Unable to load MQTT discovery")
        return False

    success = await discovery.async_start(
        hass, conf[CONF_DISCOVERY_PREFIX], hass_config,
        config_entry)  # type: bool

    return success


async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
    """Start the MQTT protocol service."""
    conf = config.get(DOMAIN)  # type: Optional[ConfigType]

    # We need this because discovery can cause components to be set up and
    # otherwise it will not load the users config.
    # This needs a better solution.
    hass.data[DATA_MQTT_HASS_CONFIG] = config

    if conf is None:
        # If we have a config entry, setup is done by that config entry.
        # If there is no config entry, this should fail.
        return bool(hass.config_entries.async_entries(DOMAIN))

    conf = dict(conf)

    if CONF_EMBEDDED in conf or CONF_BROKER not in conf:
        if (conf.get(CONF_PASSWORD) is None and
                config.get('http', {}).get('api_password') is not None):
            _LOGGER.error(
                "Starting from release 0.76, the embedded MQTT broker does not"
                " use api_password as default password anymore. Please set"
                " password configuration. See https://home-assistant.io/docs/"
                "mqtt/broker#embedded-broker for details")
            return False

        broker_config = await _async_setup_server(hass, config)

        if broker_config is None:
            _LOGGER.error("Unable to start embedded MQTT broker")
            return False

        conf.update({
            CONF_BROKER: broker_config[0],
            CONF_PORT: broker_config[1],
            CONF_USERNAME: broker_config[2],
            CONF_PASSWORD: broker_config[3],
            CONF_CERTIFICATE: broker_config[4],
            CONF_PROTOCOL: broker_config[5],
            CONF_CLIENT_KEY: None,
            CONF_CLIENT_CERT: None,
            CONF_TLS_INSECURE: None,
        })

    hass.data[DATA_MQTT_CONFIG] = conf

    # Only import if we haven't before.
    if not hass.config_entries.async_entries(DOMAIN):
        hass.async_create_task(hass.config_entries.flow.async_init(
            DOMAIN, context={'source': config_entries.SOURCE_IMPORT},
            data={}
        ))

    return True


async def async_setup_entry(hass, entry):
    """Load a config entry."""
    conf = hass.data.get(DATA_MQTT_CONFIG)

    # Config entry was created because user had configuration.yaml entry
    # They removed that, so remove entry.
    if conf is None and entry.source == config_entries.SOURCE_IMPORT:
        hass.async_create_task(
            hass.config_entries.async_remove(entry.entry_id))
        return False

    # If user didn't have configuration.yaml config, generate defaults
    if conf is None:
        conf = CONFIG_SCHEMA({
            DOMAIN: entry.data,
        })[DOMAIN]
    elif any(key in conf for key in entry.data):
        _LOGGER.warning(
            "Data in your config entry is going to override your "
            "configuration.yaml: %s", entry.data)

    conf.update(entry.data)

    broker = conf[CONF_BROKER]
    port = conf[CONF_PORT]
    client_id = conf.get(CONF_CLIENT_ID)
    keepalive = conf[CONF_KEEPALIVE]
    username = conf.get(CONF_USERNAME)
    password = conf.get(CONF_PASSWORD)
    certificate = conf.get(CONF_CERTIFICATE)
    client_key = conf.get(CONF_CLIENT_KEY)
    client_cert = conf.get(CONF_CLIENT_CERT)
    tls_insecure = conf.get(CONF_TLS_INSECURE)
    protocol = conf[CONF_PROTOCOL]

    # For cloudmqtt.com, secured connection, auto fill in certificate
    if (certificate is None and 19999 < conf[CONF_PORT] < 30000 and
            broker.endswith('.cloudmqtt.com')):
        certificate = os.path.join(
            os.path.dirname(__file__), 'addtrustexternalcaroot.crt')

    # When the certificate is set to auto, use bundled certs from requests
    elif certificate == 'auto':
        certificate = requests.certs.where()

    if CONF_WILL_MESSAGE in conf:
        will_message = Message(**conf[CONF_WILL_MESSAGE])
    else:
        will_message = None

    if CONF_BIRTH_MESSAGE in conf:
        birth_message = Message(**conf[CONF_BIRTH_MESSAGE])
    else:
        birth_message = None

    # Be able to override versions other than TLSv1.0 under Python3.6
    conf_tls_version = conf.get(CONF_TLS_VERSION)  # type: str
    if conf_tls_version == '1.2':
        tls_version = ssl.PROTOCOL_TLSv1_2
    elif conf_tls_version == '1.1':
        tls_version = ssl.PROTOCOL_TLSv1_1
    elif conf_tls_version == '1.0':
        tls_version = ssl.PROTOCOL_TLSv1
    else:
        import sys
        # Python3.6 supports automatic negotiation of highest TLS version
        if sys.hexversion >= 0x03060000:
            tls_version = ssl.PROTOCOL_TLS  # pylint: disable=no-member
        else:
            tls_version = ssl.PROTOCOL_TLSv1

    hass.data[DATA_MQTT] = MQTT(
        hass,
        broker=broker,
        port=port,
        client_id=client_id,
        keepalive=keepalive,
        username=username,
        password=password,
        certificate=certificate,
        client_key=client_key,
        client_cert=client_cert,
        tls_insecure=tls_insecure,
        protocol=protocol,
        will_message=will_message,
        birth_message=birth_message,
        tls_version=tls_version,
    )

    success = await hass.data[DATA_MQTT].async_connect()  # type: bool

    if not success:
        return False

    async def async_stop_mqtt(event: Event):
        """Stop MQTT component."""
        await hass.data[DATA_MQTT].async_disconnect()

    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, async_stop_mqtt)

    async def async_publish_service(call: ServiceCall):
        """Handle MQTT publish service calls."""
        msg_topic = call.data[ATTR_TOPIC]  # type: str
        payload = call.data.get(ATTR_PAYLOAD)
        payload_template = call.data.get(ATTR_PAYLOAD_TEMPLATE)
        qos = call.data[ATTR_QOS]  # type: int
        retain = call.data[ATTR_RETAIN]  # type: bool
        if payload_template is not None:
            try:
                payload = \
                    template.Template(payload_template, hass).async_render()
            except template.jinja2.TemplateError as exc:
                _LOGGER.error(
                    "Unable to publish to %s: rendering payload template of "
                    "%s failed because %s",
                    msg_topic, payload_template, exc)
                return

        await hass.data[DATA_MQTT].async_publish(
            msg_topic, payload, qos, retain)

    hass.services.async_register(
        DOMAIN, SERVICE_PUBLISH, async_publish_service,
        schema=MQTT_PUBLISH_SCHEMA)

    if conf.get(CONF_DISCOVERY):
        await _async_setup_discovery(
            hass, conf, hass.data[DATA_MQTT_HASS_CONFIG], entry)

    return True


@attr.s(slots=True, frozen=True)
class Subscription:
    """Class to hold data about an active subscription."""

    topic = attr.ib(type=str)
    callback = attr.ib(type=MessageCallbackType)
    qos = attr.ib(type=int, default=0)
    encoding = attr.ib(type=str, default='utf-8')


@attr.s(slots=True, frozen=True)
class Message:
    """MQTT Message."""

    topic = attr.ib(type=str)
    payload = attr.ib(type=PublishPayloadType)
    qos = attr.ib(type=int, default=0)
    retain = attr.ib(type=bool, default=False)


class MQTT:
    """Home Assistant MQTT client."""

    def __init__(self, hass: HomeAssistantType, broker: str, port: int,
                 client_id: Optional[str], keepalive: Optional[int],
                 username: Optional[str], password: Optional[str],
                 certificate: Optional[str], client_key: Optional[str],
                 client_cert: Optional[str], tls_insecure: Optional[bool],
                 protocol: Optional[str], will_message: Optional[Message],
                 birth_message: Optional[Message],
                 tls_version: Optional[int]) -> None:
        """Initialize Home Assistant MQTT client."""
        import paho.mqtt.client as mqtt

        self.hass = hass
        self.broker = broker
        self.port = port
        self.keepalive = keepalive
        self.subscriptions = []  # type: List[Subscription]
        self.birth_message = birth_message
        self._mqttc = None  # type: mqtt.Client
        self._paho_lock = asyncio.Lock(loop=hass.loop)

        if protocol == PROTOCOL_31:
            proto = mqtt.MQTTv31  # type: int
        else:
            proto = mqtt.MQTTv311

        if client_id is None:
            self._mqttc = mqtt.Client(protocol=proto)
        else:
            self._mqttc = mqtt.Client(client_id, protocol=proto)

        if username is not None:
            self._mqttc.username_pw_set(username, password)

        if certificate is not None:
            self._mqttc.tls_set(
                certificate, certfile=client_cert,
                keyfile=client_key, tls_version=tls_version)

            if tls_insecure is not None:
                self._mqttc.tls_insecure_set(tls_insecure)

        self._mqttc.on_connect = self._mqtt_on_connect
        self._mqttc.on_disconnect = self._mqtt_on_disconnect
        self._mqttc.on_message = self._mqtt_on_message

        if will_message is not None:
            self._mqttc.will_set(*attr.astuple(will_message))

    async def async_publish(self, topic: str, payload: PublishPayloadType,
                            qos: int, retain: bool) -> None:
        """Publish a MQTT message.

        This method must be run in the event loop and returns a coroutine.
        """
        async with self._paho_lock:
            _LOGGER.debug("Transmitting message on %s: %s", topic, payload)
            await self.hass.async_add_job(
                self._mqttc.publish, topic, payload, qos, retain)

    async def async_connect(self) -> bool:
        """Connect to the host. Does process messages yet.

        This method is a coroutine.
        """
        result = None  # type: int
        try:
            result = await self.hass.async_add_job(
                self._mqttc.connect, self.broker, self.port, self.keepalive)
        except OSError as err:
            _LOGGER.error("Failed to connect due to exception: %s", err)
            return False

        if result != 0:
            import paho.mqtt.client as mqtt
            _LOGGER.error("Failed to connect: %s", mqtt.error_string(result))
            return False

        self._mqttc.loop_start()
        return True

    @callback
    def async_disconnect(self):
        """Stop the MQTT client.

        This method must be run in the event loop and returns a coroutine.
        """
        def stop():
            """Stop the MQTT client."""
            self._mqttc.disconnect()
            self._mqttc.loop_stop()

        return self.hass.async_add_job(stop)

    async def async_subscribe(self, topic: str,
                              msg_callback: MessageCallbackType,
                              qos: int, encoding: str) -> Callable[[], None]:
        """Set up a subscription to a topic with the provided qos.

        This method is a coroutine.
        """
        if not isinstance(topic, str):
            raise HomeAssistantError("Topic needs to be a string!")

        subscription = Subscription(topic, msg_callback, qos, encoding)
        self.subscriptions.append(subscription)

        await self._async_perform_subscription(topic, qos)

        @callback
        def async_remove() -> None:
            """Remove subscription."""
            if subscription not in self.subscriptions:
                raise HomeAssistantError("Can't remove subscription twice")
            self.subscriptions.remove(subscription)

            if any(other.topic == topic for other in self.subscriptions):
                # Other subscriptions on topic remaining - don't unsubscribe.
                return
            self.hass.async_create_task(self._async_unsubscribe(topic))

        return async_remove

    async def _async_unsubscribe(self, topic: str) -> None:
        """Unsubscribe from a topic.

        This method is a coroutine.
        """
        async with self._paho_lock:
            result = None  # type: int
            result, _ = await self.hass.async_add_job(
                self._mqttc.unsubscribe, topic)
            _raise_on_error(result)

    async def _async_perform_subscription(self, topic: str, qos: int) -> None:
        """Perform a paho-mqtt subscription."""
        _LOGGER.debug("Subscribing to %s", topic)

        async with self._paho_lock:
            result = None  # type: int
            result, _ = await self.hass.async_add_job(
                self._mqttc.subscribe, topic, qos)
            _raise_on_error(result)

    def _mqtt_on_connect(
            self, _mqttc, _userdata, _flags, result_code: int) -> None:
        """On connect callback.

        Resubscribe to all topics we were subscribed to and publish birth
        message.
        """
        import paho.mqtt.client as mqtt

        if result_code != mqtt.CONNACK_ACCEPTED:
            _LOGGER.error("Unable to connect to the MQTT broker: %s",
                          mqtt.connack_string(result_code))
            self._mqttc.disconnect()
            return

        # Group subscriptions to only re-subscribe once for each topic.
        keyfunc = attrgetter('topic')
        for topic, subs in groupby(sorted(self.subscriptions, key=keyfunc),
                                   keyfunc):
            # Re-subscribe with the highest requested qos
            max_qos = max(subscription.qos for subscription in subs)
            self.hass.add_job(self._async_perform_subscription, topic, max_qos)

        if self.birth_message:
            self.hass.add_job(
                self.async_publish(*attr.astuple(self.birth_message)))

    def _mqtt_on_message(self, _mqttc, _userdata, msg) -> None:
        """Message received callback."""
        self.hass.add_job(self._mqtt_handle_message, msg)

    @callback
    def _mqtt_handle_message(self, msg) -> None:
        _LOGGER.debug("Received message on %s: %s", msg.topic, msg.payload)

        for subscription in self.subscriptions:
            if not _match_topic(subscription.topic, msg.topic):
                continue

            payload = msg.payload  # type: SubscribePayloadType
            if subscription.encoding is not None:
                try:
                    payload = msg.payload.decode(subscription.encoding)
                except (AttributeError, UnicodeDecodeError):
                    _LOGGER.warning(
                        "Can't decode payload %s on %s with encoding %s",
                        msg.payload, msg.topic, subscription.encoding)
                    continue

            self.hass.async_run_job(
                subscription.callback, msg.topic, payload, msg.qos)

    def _mqtt_on_disconnect(self, _mqttc, _userdata, result_code: int) -> None:
        """Disconnected callback."""
        # When disconnected because of calling disconnect()
        if result_code == 0:
            return

        tries = 0

        while True:
            try:
                if self._mqttc.reconnect() == 0:
                    _LOGGER.info("Successfully reconnected to the MQTT server")
                    break
            except socket.error:
                pass

            wait_time = min(2**tries, MAX_RECONNECT_WAIT)
            _LOGGER.warning(
                "Disconnected from MQTT (%s). Trying to reconnect in %s s",
                result_code, wait_time)
            # It is ok to sleep here as we are in the MQTT thread.
            time.sleep(wait_time)
            tries += 1


def _raise_on_error(result_code: int) -> None:
    """Raise error if error result."""
    if result_code != 0:
        import paho.mqtt.client as mqtt

        raise HomeAssistantError(
            'Error talking to MQTT: {}'.format(mqtt.error_string(result_code)))


def _match_topic(subscription: str, topic: str) -> bool:
    """Test if topic matches subscription."""
    from paho.mqtt.matcher import MQTTMatcher
    matcher = MQTTMatcher()
    matcher[subscription] = True
    try:
        next(matcher.iter_match(topic))
        return True
    except StopIteration:
        return False


class MqttAttributes(Entity):
    """Mixin used for platforms that support JSON attributes."""

    def __init__(self, config: dict) -> None:
        """Initialize the JSON attributes mixin."""
        self._attributes = None
        self._attributes_sub_state = None
        self._attributes_config = config

    async def async_added_to_hass(self) -> None:
        """Subscribe MQTT events.

        This method must be run in the event loop and returns a coroutine.
        """
        await super().async_added_to_hass()
        await self._attributes_subscribe_topics()

    async def attributes_discovery_update(self, config: dict):
        """Handle updated discovery message."""
        self._attributes_config = config
        await self._attributes_subscribe_topics()

    async def _attributes_subscribe_topics(self):
        """(Re)Subscribe to topics."""
        from .subscription import async_subscribe_topics

        @callback
        def attributes_message_received(topic: str,
                                        payload: SubscribePayloadType,
                                        qos: int) -> None:
            try:
                json_dict = json.loads(payload)
                if isinstance(json_dict, dict):
                    self._attributes = json_dict
                    self.async_schedule_update_ha_state()
                else:
                    _LOGGER.warning("JSON result was not a dictionary")
                    self._attributes = None
            except ValueError:
                _LOGGER.warning("Erroneous JSON: %s", payload)
                self._attributes = None

        self._attributes_sub_state = await async_subscribe_topics(
            self.hass, self._attributes_sub_state,
            {CONF_JSON_ATTRS_TOPIC: {
                'topic': self._attributes_config.get(CONF_JSON_ATTRS_TOPIC),
                'msg_callback': attributes_message_received,
                'qos': self._attributes_config.get(CONF_QOS)}})

    async def async_will_remove_from_hass(self):
        """Unsubscribe when removed."""
        from .subscription import async_unsubscribe_topics
        self._attributes_sub_state = await async_unsubscribe_topics(
            self.hass, self._attributes_sub_state)

    @property
    def device_state_attributes(self):
        """Return the state attributes."""
        return self._attributes


class MqttAvailability(Entity):
    """Mixin used for platforms that report availability."""

    def __init__(self, config: dict) -> None:
        """Initialize the availability mixin."""
        self._availability_sub_state = None
        self._available = False  # type: bool

        self._avail_config = config

    async def async_added_to_hass(self) -> None:
        """Subscribe MQTT events.

        This method must be run in the event loop and returns a coroutine.
        """
        await super().async_added_to_hass()
        await self._availability_subscribe_topics()

    async def availability_discovery_update(self, config: dict):
        """Handle updated discovery message."""
        self._avail_config = config
        await self._availability_subscribe_topics()

    async def _availability_subscribe_topics(self):
        """(Re)Subscribe to topics."""
        from .subscription import async_subscribe_topics

        @callback
        def availability_message_received(topic: str,
                                          payload: SubscribePayloadType,
                                          qos: int) -> None:
            """Handle a new received MQTT availability message."""
            if payload == self._avail_config[CONF_PAYLOAD_AVAILABLE]:
                self._available = True
            elif payload == self._avail_config[CONF_PAYLOAD_NOT_AVAILABLE]:
                self._available = False

            self.async_schedule_update_ha_state()

        self._availability_sub_state = await async_subscribe_topics(
            self.hass, self._availability_sub_state,
            {'availability_topic': {
                'topic': self._avail_config.get(CONF_AVAILABILITY_TOPIC),
                'msg_callback': availability_message_received,
                'qos': self._avail_config[CONF_QOS]}})

    async def async_will_remove_from_hass(self):
        """Unsubscribe when removed."""
        from .subscription import async_unsubscribe_topics
        self._availability_sub_state = await async_unsubscribe_topics(
            self.hass, self._availability_sub_state)

    @property
    def available(self) -> bool:
        """Return if the device is available."""
        availability_topic = self._avail_config.get(CONF_AVAILABILITY_TOPIC)
        return availability_topic is None or self._available


class MqttDiscoveryUpdate(Entity):
    """Mixin used to handle updated discovery message."""

    def __init__(self, discovery_hash, discovery_update=None) -> None:
        """Initialize the discovery update mixin."""
        self._discovery_hash = discovery_hash
        self._discovery_update = discovery_update
        self._remove_signal = None

    async def async_added_to_hass(self) -> None:
        """Subscribe to discovery updates."""
        await super().async_added_to_hass()

        from homeassistant.helpers.dispatcher import async_dispatcher_connect
        from homeassistant.components.mqtt.discovery import (
            MQTT_DISCOVERY_UPDATED, clear_discovery_hash)

        @callback
        def discovery_callback(payload):
            """Handle discovery update."""
            _LOGGER.info("Got update for entity with hash: %s '%s'",
                         self._discovery_hash, payload)
            if not payload:
                # Empty payload: Remove component
                _LOGGER.info("Removing component: %s", self.entity_id)
                self.hass.async_create_task(self.async_remove())
                clear_discovery_hash(self.hass, self._discovery_hash)
                self._remove_signal()
            elif self._discovery_update:
                # Non-empty payload: Notify component
                _LOGGER.info("Updating component: %s", self.entity_id)
                payload.pop(ATTR_DISCOVERY_HASH)
                self.hass.async_create_task(self._discovery_update(payload))

        if self._discovery_hash:
            self._remove_signal = async_dispatcher_connect(
                self.hass,
                MQTT_DISCOVERY_UPDATED.format(self._discovery_hash),
                discovery_callback)


class MqttEntityDeviceInfo(Entity):
    """Mixin used for mqtt platforms that support the device registry."""

    def __init__(self, device_config: Optional[ConfigType],
                 config_entry=None) -> None:
        """Initialize the device mixin."""
        self._device_config = device_config
        self._config_entry = config_entry

    async def device_info_discovery_update(self, config: dict):
        """Handle updated discovery message."""
        self._device_config = config.get(CONF_DEVICE)
        device_registry = await \
            self.hass.helpers.device_registry.async_get_registry()
        config_entry_id = self._config_entry.entry_id
        device_info = self.device_info

        if config_entry_id is not None and device_info is not None:
            device_info['config_entry_id'] = config_entry_id
            device_registry.async_get_or_create(**device_info)

    @property
    def device_info(self):
        """Return a device description for device registry."""
        if not self._device_config:
            return None

        info = {
            'identifiers': {
                (DOMAIN, id_)
                for id_ in self._device_config[CONF_IDENTIFIERS]
            },
            'connections': {
                tuple(x) for x in self._device_config[CONF_CONNECTIONS]
            }
        }

        if CONF_MANUFACTURER in self._device_config:
            info['manufacturer'] = self._device_config[CONF_MANUFACTURER]

        if CONF_MODEL in self._device_config:
            info['model'] = self._device_config[CONF_MODEL]

        if CONF_NAME in self._device_config:
            info['name'] = self._device_config[CONF_NAME]

        if CONF_SW_VERSION in self._device_config:
            info['sw_version'] = self._device_config[CONF_SW_VERSION]

        if CONF_VIA_HUB in self._device_config:
            info['via_hub'] = (DOMAIN, self._device_config[CONF_VIA_HUB])

        return info
