# Copyright © The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.
"""Tests for ServerNotifications."""

import asyncio
import logging
import signal
from typing import Any, override
from unittest import mock
from unittest.mock import MagicMock

import tenacity

from debusine.client.debusine import Debusine, ServerNotificationDecodeError
from debusine.client.server_notifications import ServerNotifications, logger
from debusine.test import TestCase
from debusine.worker.config import ConfigHandler
from debusine.worker.tests import server
from debusine.worker.tests.server import ServerConfig


class MainLoopKeepsRetrying(Exception):
    """Exception raised when tenacity does not stop retrying."""


class ServerNotificationsTest(TestCase, server.DebusineAioHTTPTestCase):
    """Test :py:class:`ServerNotifications`."""

    @override
    async def setUpAsync(self) -> None:
        await super().setUpAsync()
        self.api_url = str(self.server.make_url('')) + "/api"
        self.temp_config_directory = self.create_temp_config_directory(
            {'General': {'api-url': self.api_url}}
        )
        self.config = ConfigHandler(
            directories=[self.temp_config_directory.as_posix()]
        )
        self.config['General']['log-file'] = str(
            self.temp_config_directory / 'worker.log'
        )

        # Set a valid token in the test server and client configuration
        token = '6c931875627131b5135b7de3371c44'
        self.server_config.registered_token = token
        self.config.write_token(token)

        self.default_sigint_handler = signal.getsignal(signal.SIGINT)
        self.default_sigterm_handler = signal.getsignal(signal.SIGTERM)

        self._server_notifications: ServerNotifications | None = None

    @override
    def tearDown(self) -> None:
        super().tearDown()
        # Restore signal handlers. Cli.execute() changes them
        signal.signal(signal.SIGINT, self.default_sigint_handler)
        signal.signal(signal.SIGTERM, self.default_sigterm_handler)

    @override
    async def tearDownAsync(self) -> None:
        if self._server_notifications:
            await self._server_notifications.close()
        await super().tearDownAsync()

    @property
    def server_notifications(self) -> ServerNotifications:
        """Return the ServerNotifications being tested."""
        if self._server_notifications is None:
            debusine_client = Debusine(
                self.config.api_url,
                self.config.token,
                timeout=self.config.http_timeout,
                logger=logging.getLogger(),
                scope=None,
            )
            self._server_notifications = ServerNotifications(
                debusine_client, endpoint="1.0/worker/connect/"
            )
            self._server_notification_asyncio_sleep = self.enterContext(
                mock.patch.object(
                    self.server_notifications,
                    "_asyncio_sleep",
                    side_effect=[None] * 5 + [MainLoopKeepsRetrying()],
                )
            )
        return self._server_notifications

    def mock_wait_for_messages(self) -> asyncio.Future[None]:
        """
        Mock the message reading function with a wait for a future.

        :returns: the future that blocks the mocked main loop: set a result to
          exit the mocked main loop, set an exception to make it raise, cancel
          it to cancel the main loop
        """
        future = asyncio.get_running_loop().create_future()

        async def mock_wait_for_messages() -> None:
            await future

        self.enterContext(
            mock.patch.object(
                self.server_notifications,
                "_wait_for_messages",
                mock_wait_for_messages,
            )
        )
        return future

    async def test_token_required(self) -> None:
        client = Debusine(
            self.config.api_url,
            None,
            timeout=self.config.http_timeout,
            logger=logging.getLogger(),
            scope=None,
        )
        with self.assertRaisesRegex(
            ValueError,
            "Debusine client needs to be configured with an API token",
        ):
            ServerNotifications(client, endpoint="1.0/worker/connect/")

    async def test_asyncio_sleep(self) -> None:
        # Reinstantiate ServerNotifications to get a version with an unpatched
        # _asyncio_sleep method
        server_notifications = ServerNotifications(
            self.server_notifications.client,
            endpoint=self.server_notifications.endpoint,
        )
        with mock.patch("asyncio.sleep") as asleep:
            await server_notifications._asyncio_sleep(4.2)
        asleep.assert_called_once_with(4.2)

    async def test_connection_success(self) -> None:
        # Do not try to reconnect when the server quits
        self.server_notifications._shutting_down = True
        await self.server_notifications.main_loop()
        assert self.server_latest_request is not None
        self.assertEqual(
            self.server_latest_request.path, '/api/ws/1.0/worker/connect/'
        )
        self.assertTrue(self.server_latest_request.is_authenticated)

    async def test_invalid_token(self) -> None:
        self.registered_token = 'f924ba0099588bf6'
        self.config.write_token('e3c0284')
        # Do not try to reconnect when the server quits
        self.server_notifications._shutting_down = True
        await self.server_notifications.main_loop()
        assert self.server_latest_request is not None
        self.assertEqual(
            self.server_latest_request.path, '/api/ws/1.0/worker/connect/'
        )
        self.assertFalse(self.server_latest_request.is_authenticated)

    async def test_connection_reset(self) -> None:
        self.enterContext(
            mock.patch.object(
                self.server_notifications,
                "_raise_try_again",
                side_effect=ConnectionResetError,
            )
        )

        log_message = 'ConnectionResetError, will reconnect'
        with (
            self.assertLogsContains(
                log_message,
                expected_count=6,
                level=logging.DEBUG,
                logger=logger,
            ),
            self.assertRaises(MainLoopKeepsRetrying),
        ):
            await self.server_notifications.main_loop()

    async def test_connection_refused(self) -> None:
        await self.server.close()

        log_message = 'Error connecting to'
        with (
            self.assertLogsContains(
                log_message, expected_count=6, logger=logger
            ),
            self.assertRaises(MainLoopKeepsRetrying),
        ):
            await self.server_notifications.main_loop()

    async def test_remote_close_exits_main_loop(self) -> None:
        self.server_config.connect_response_bodies = [
            ServerConfig.RESPONSES["request_dynamic_metadata"]
        ]

        async def call_close(*args: Any) -> None:
            await self.server_notifications.close()

        self.enterContext(
            mock.patch.object(
                self.server_notifications.message_queue,
                "put",
                side_effect=call_close,
            )
        )
        with (
            mock.patch.object(
                self.server_notifications, "_raise_try_again"
            ) as try_again_mock,
        ):
            await self.server_notifications.main_loop()
        try_again_mock.assert_not_called()

    async def test_token_disabled_try_again(self) -> None:
        assert self.config.token
        # The client will reconnect, so the server will keep resending this
        # sequence of messages every time
        self.server_config.connect_response_bodies = [
            ServerConfig.RESPONSES["token_disabled_error"],
        ]
        with (
            self.assertLogsContains(
                "The token is disabled. Will try again",
                expected_count=6,
                logger=logger,
            ),
            self.assertRaises(MainLoopKeepsRetrying),
        ):
            await self.server_notifications.main_loop()

    async def test_decode_ws_message_fails(self) -> None:
        self.enterContext(
            mock.patch.object(
                self.server_notifications.client,
                "decode_ws_message",
                side_effect=ServerNotificationDecodeError("expected error"),
            )
        )
        self.server_config.connect_response_bodies = [
            ServerConfig.RESPONSES["request_dynamic_metadata"],
            ServerConfig.RESPONSES["request_dynamic_metadata"],
        ]

        self.server_notifications._shutting_down = True
        with self.assertLogs(level=logging.WARNING, logger=logger) as log:
            await self.server_notifications.main_loop()

        # Messages were ignored
        self.assertEqual(self.server_notifications.message_queue.qsize(), 0)

        msg = (
            f"WARNING:{logger.name}:Unhandled websocket notification: "
            f"expected error"
        )
        self.assertEqual(log.output, [msg, msg])

    async def test_messages_enqueued(self) -> None:
        RESPONSES = ServerConfig.RESPONSES
        self.server_config.connect_response_bodies = [
            RESPONSES["connected"],
            RESPONSES["request_dynamic_metadata"],
            RESPONSES["work_request_available"],
            RESPONSES["request_dynamic_metadata"],
        ]
        # Quit the main loop when the server stops sending
        self.server_notifications._shutting_down = True
        await self.server_notifications.main_loop()

        self.assertEqual(self.server_notifications.message_queue.qsize(), 4)
        payload = self.server_notifications.message_queue.get_nowait()
        self.assertEqual(payload, RESPONSES["connected"])
        payload = self.server_notifications.message_queue.get_nowait()
        self.assertEqual(payload, RESPONSES["request_dynamic_metadata"])
        payload = self.server_notifications.message_queue.get_nowait()
        self.assertEqual(payload, RESPONSES["work_request_available"])
        payload = self.server_notifications.message_queue.get_nowait()
        self.assertEqual(payload, RESPONSES["request_dynamic_metadata"])

    def test_raise_try_again(self) -> None:
        self.assertRaises(
            tenacity.TryAgain, self.server_notifications._raise_try_again
        )

    def test_reset_retry_counter(self) -> None:
        retry_state = MagicMock()
        retry_state.attempt_number = 5

        self.server_notifications._reset_retry_attempt_number = True
        self.server_notifications._reset_retry_counter(retry_state)
        self.assertEqual(retry_state.attempt_number, 0)
        self.assertFalse(self.server_notifications._reset_retry_attempt_number)

    async def test_open(self) -> None:
        main_loop_future = self.mock_wait_for_messages()
        self.assertIsNone(self.server_notifications.main_loop_task)
        await self.server_notifications.open()
        self.assertIsNotNone(self.server_notifications.main_loop_task)
        main_loop_future.set_result(None)
        assert self.server_notifications.main_loop_task is not None
        await self.server_notifications.main_loop_task

    async def test_context_manager(self) -> None:
        self.mock_wait_for_messages()
        self.assertIsNone(self.server_notifications.main_loop_task)
        async with self.server_notifications:
            self.assertIsNotNone(self.server_notifications.main_loop_task)

    async def test_generator(self) -> None:
        RESPONSES = ServerConfig.RESPONSES
        expected = [
            RESPONSES["connected"],
            RESPONSES["request_dynamic_metadata"],
            RESPONSES["work_request_available"],
            RESPONSES["request_dynamic_metadata"],
        ]
        for payload in expected:
            await self.server_notifications.message_queue.put(payload)

        main_loop_future = self.mock_wait_for_messages()

        actual: list[dict[str, Any]] = []
        async with self.server_notifications:
            async for payload in self.server_notifications.messages():
                actual.append(payload)
                if len(actual) == len(expected):
                    main_loop_future.set_result(None)

        assert actual == expected

    async def test_close_when_iterating(self) -> None:
        RESPONSES = ServerConfig.RESPONSES
        await self.server_notifications.message_queue.put(
            RESPONSES["connected"]
        )
        await self.server_notifications.message_queue.put(
            RESPONSES["request_dynamic_metadata"]
        )

        self.mock_wait_for_messages()

        actual: list[dict[str, Any]] = []
        async with self.server_notifications:
            async for payload in self.server_notifications.messages():
                actual.append(payload)
                await self.server_notifications.close()

        assert actual == [RESPONSES["connected"]]
