import asyncio
import logging
from importlib.metadata import EntryPoint
from logging.config import dictConfig
from unittest.mock import patch

import pytest

from amqtt.broker import Broker
from amqtt.client import MQTTClient
from amqtt.mqtt.constants import QOS_0
from tests.asserts import does_not_warn

dictConfig({
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {
        'verbose': {
            'format': '%(asctime)s [%(levelname)s] %(name)s: %(message)s'
        }
    },
    'handlers': {
        'console': {
            'class': 'logging.StreamHandler',
            'level': 'DEBUG',
            'formatter': 'verbose',
        }
    },
    'loggers': {
        'transitions': {
            'level': 'WARNING',
        }
    }
})

# logging.basicConfig(level=logging.DEBUG, format=formatter)

logger = logging.getLogger(__name__)


all_sys_topics = [
    '$SYS/broker/version',
    '$SYS/broker/load/bytes/received',
    '$SYS/broker/load/bytes/sent',
    '$SYS/broker/messages/received',
    '$SYS/broker/messages/sent',
    '$SYS/broker/time',
    '$SYS/broker/uptime',
    '$SYS/broker/uptime/formatted',
    '$SYS/broker/clients/connected',
    '$SYS/broker/clients/disconnected',
    '$SYS/broker/clients/maximum',
    '$SYS/broker/clients/total',
    '$SYS/broker/messages/inflight',
    '$SYS/broker/messages/inflight/in',
    '$SYS/broker/messages/inflight/out',
    '$SYS/broker/messages/inflight/stored',
    '$SYS/broker/messages/publish/received',
    '$SYS/broker/messages/publish/sent',
    '$SYS/broker/messages/retained/count',
    '$SYS/broker/messages/subscriptions/count',
    '$SYS/broker/heap/size',
    '$SYS/broker/heap/maximum',
    '$SYS/broker/cpu/percent',
    '$SYS/broker/cpu/maximum',
]


async def _collect_expected_sys_topics(client: MQTTClient, sys_topic_flags: dict[str, bool]) -> int:
    loop = asyncio.get_running_loop()
    deadline = loop.time() + 5
    sys_msg_count = 0

    while not all(sys_topic_flags.values()):
        remaining = deadline - loop.time()
        if remaining <= 0:
            break

        try:
            message = await client.deliver_message(timeout_duration=min(remaining, 1))
        except asyncio.TimeoutError:
            logger.debug(f"TimeoutError after {sys_msg_count} messages")
            continue

        if message and message.topic.startswith("$SYS/"):
            sys_msg_count += 1
            assert message.topic in sys_topic_flags
            sys_topic_flags[message.topic] = True

    return sys_msg_count


def _cancel_sys_broadcast(broker: Broker) -> None:
    sys_plugin = broker.plugins_manager.get_plugin("BrokerSysPlugin")
    sys_handle = getattr(sys_plugin, "_sys_handle", None)
    if sys_handle:
        sys_handle.cancel()


def _broker_uri(broker: Broker) -> str:
    sockname = broker._servers["default"].instance.sockets[0].getsockname()
    return f"mqtt://127.0.0.1:{sockname[1]}/"


async def _disconnect_client(client: MQTTClient) -> None:
    if client.session and client.session.transitions.is_connected():
        await client.disconnect()


async def _shutdown_broker(broker: Broker) -> None:
    if broker.transitions.is_started():
        await broker.shutdown()



# test broker sys
@pytest.mark.asyncio
async def test_broker_sys_plugin_deprecated_config() -> None:

    sys_topic_flags = {sys_topic:False for sys_topic in all_sys_topics}

    class MockEntryPoints:

        def select(self, group) -> list[EntryPoint]:
            match group:
                case 'tests.mock_plugins':
                    return [
                            EntryPoint(name='broker_sys', group='tests.mock_plugins', value='amqtt.plugins.sys.broker:BrokerSysPlugin'),
                            EntryPoint(name='auth_anonymous', group='test.mock_plugins', value='amqtt.plugins.authentication:AnonymousAuthPlugin'),
                        ]
                case _:
                    return list()


    with patch("amqtt.plugins.manager.entry_points", side_effect=MockEntryPoints) as mocked_mqtt_publish:

        config = {
            "listeners": {
                "default": {"type": "tcp", "bind": "127.0.0.1:0", "max_connections": 10},
            },
            'sys_interval': 1,
            'auth': {
                'allow_anonymous': True
            }
        }

        broker = Broker(plugin_namespace='tests.mock_plugins', config=config)
        client = MQTTClient()
        sys_msg_count = 0
        try:
            await broker.start()
            await client.connect(_broker_uri(broker))
            await client.subscribe([("$SYS/#", QOS_0),])
            await client.publish('test/topic', b'my test message')
            sys_msg_count = await _collect_expected_sys_topics(client, sys_topic_flags)
        finally:
            _cancel_sys_broadcast(broker)
            await _disconnect_client(client)
            await _shutdown_broker(broker)

        assert sys_msg_count > 1

        assert all(sys_topic_flags.values()), f'topic not received: {[ topic for topic, flag in sys_topic_flags.items() if not flag ]}'


@pytest.mark.asyncio
async def test_broker_sys_plugin_config() -> None:

    sys_topic_flags = {sys_topic:False for sys_topic in all_sys_topics}

    config = {
        "listeners": {
            "default": {"type": "tcp", "bind": "127.0.0.1:0", "max_connections": 10},
        },
        'plugins': [
            {'amqtt.plugins.authentication.AnonymousAuthPlugin': {'allow_anonymous': True}},
            {'amqtt.plugins.sys.broker.BrokerSysPlugin': {'sys_interval': 1}},
        ]
    }


    broker = Broker(plugin_namespace='tests.mock_plugins', config=config)
    client = MQTTClient()
    sys_msg_count = 0
    try:
        await broker.start()
        await client.connect(_broker_uri(broker))
        await client.subscribe([("$SYS/#", QOS_0), ])
        await client.publish('test/topic', b'my test message')
        sys_msg_count = await _collect_expected_sys_topics(client, sys_topic_flags)
    finally:
        _cancel_sys_broadcast(broker)
        await _disconnect_client(client)
        await _shutdown_broker(broker)

    assert sys_msg_count > 1

    assert all(
        sys_topic_flags.values()), f'topic not received: {[topic for topic, flag in sys_topic_flags.items() if not flag]}'


@pytest.mark.asyncio
async def test_broker_sys_plugin_without_interval_set(caplog) -> None:

    sys_topic_flags = {sys_topic:False for sys_topic in all_sys_topics}

    config = {
        "listeners": {
            "default": {"type": "tcp", "bind": "127.0.0.1:0", "max_connections": 10},
        },
        'plugins': [
            {'amqtt.plugins.authentication.AnonymousAuthPlugin': {'allow_anonymous': True}},
            {'amqtt.plugins.sys.broker.BrokerSysPlugin': {}},  # not set sys_interval, gets default of 20
        ]
    }

    broker = Broker(plugin_namespace='tests.mock_plugins', config=config)

    with caplog.at_level(logging.WARNING):
        # if sys_interval is set to `None`, the broker will throw an error that 'NoneType'
        #    which won't t match expected type 'int'
        # to reproduce #298, need to set it to `None` after the plugin is loaded but before `post_start`
        broker.plugins_manager.get_plugin('BrokerSysPlugin').context.config.sys_interval = None
        await broker.start()
        await asyncio.sleep(0.1)
        assert "'sys_interval' key is not set or is None" in caplog.text


    client = MQTTClient()
    sys_msg_count = 0
    try:
        await client.connect(_broker_uri(broker))
        await client.subscribe([("$SYS/#", QOS_0), ])
        await client.publish('test/topic', b'my test message')
        await asyncio.sleep(2)
        try:
            while sys_msg_count < 30:
                message = await client.deliver_message(timeout_duration=1)
                if '$SYS' in message.topic:
                    sys_msg_count += 1
                    assert message.topic in sys_topic_flags
                    sys_topic_flags[message.topic] = True

        except asyncio.TimeoutError:
            logger.debug(f"TimeoutError after {sys_msg_count} messages")
    finally:
        _cancel_sys_broadcast(broker)
        await _disconnect_client(client)
        await _shutdown_broker(broker)

    # the only $SYS message received should be the retained broker version
    assert sys_msg_count == 1


@pytest.mark.asyncio
async def test_broker_sys_psutil_deprecation_warning() -> None:


    with_sys_config = {
        "listeners": {
            "default": {"type": "tcp", "bind": "127.0.0.1:1883", "max_connections": 10},
        },
        'plugins': [
            {'amqtt.plugins.authentication.AnonymousAuthPlugin': {'allow_anonymous': True}},
            {'amqtt.plugins.sys.broker.BrokerSysPlugin': {'sys_interval': 1}},
        ]
    }

    no_sys_config = {
        "listeners": {
            "default": {"type": "tcp", "bind": "127.0.0.1:1883", "max_connections": 10},
        },
        'plugins': [
            {'amqtt.plugins.authentication.AnonymousAuthPlugin': {'allow_anonymous': True}},
        ]
    }


    with pytest.warns(DeprecationWarning):
        _ = Broker(plugin_namespace='tests.mock_plugins', config=with_sys_config)

    with does_not_warn():
        _ = Broker(plugin_namespace='tests.mock_plugins', config=no_sys_config)
