# 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.
"""Receive server notifications."""

import asyncio
import logging
import types
from collections.abc import AsyncGenerator
from typing import Any, NoReturn, Self

import aiohttp
import tenacity

from debusine.client.debusine import Debusine, ServerNotificationDecodeError
from debusine.client.exceptions import TokenDisabledError
from debusine.utils import get_user_agent

logger = logging.getLogger(__name__)


class ServerNotifications:
    """Read asynchronous notification from a remote Debusine server."""

    def __init__(self, client: Debusine, endpoint: str) -> None:
        """
        Initialize ServerNotifications.

        :param config: worker configuration
        :param client: scope-less Debusine client
        """
        self.client = client
        if self.client.token is None:
            raise ValueError(
                "Debusine client needs to be configured with an API token"
            )
        self.endpoint = endpoint
        self.message_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        self.client_session: aiohttp.ClientSession | None = None
        self.debusine_ws_url = self.client.base_api_url.replace('http', 'ws', 1)
        self.worker_connect_ws_url = f"{self.debusine_ws_url}/ws/{endpoint}"
        # Set to True if the retry_attempt_number should be reset
        self._reset_retry_attempt_number: bool = False

        # Running main loop
        self.main_loop_task: asyncio.Task[None] | None = None

        # Set to True to stop the main loop from looping
        self._shutting_down: bool = False

    @staticmethod
    async def _asyncio_sleep(delay: float) -> None:
        """Sleep asynchronously (mocked in tests)."""
        await asyncio.sleep(delay)

    def _tenacity_before_hook(
        self, retry_state: tenacity.RetryCallState
    ) -> None:
        self._reset_retry_counter(retry_state)
        self._log_trying_to_reconnect(retry_state)

    def _reset_retry_counter(
        self, retry_state: tenacity.RetryCallState
    ) -> None:
        if self._reset_retry_attempt_number:
            retry_state.attempt_number = 0
            self._reset_retry_attempt_number = False

    @staticmethod
    def _log_trying_to_reconnect(retry_state: tenacity.RetryCallState) -> None:
        if retry_state.attempt_number > 1:
            logger.debug(
                'Trying to reconnect, attempt %d', retry_state.attempt_number
            )

    def _raise_try_again(self) -> NoReturn:
        raise tenacity.TryAgain()

    async def open(self) -> None:
        """Start the reader."""
        if self.main_loop_task is not None:
            raise AssertionError("Reader task already started")
        self.main_loop_task = asyncio.create_task(self.main_loop())

    async def close(self) -> None:
        """Close the message reader."""
        self._shutting_down = True
        if self.client_session is not None:
            await self.client_session.close()

        if self.main_loop_task is not None:
            self.main_loop_task.cancel()
            # Await so that we do not hide potential exceptions raised in
            # main_loop
            try:
                await self.main_loop_task
            except asyncio.CancelledError:
                pass
            self.main_loop_task = None

    async def __aenter__(self) -> Self:
        """Start the reader task."""
        await self.open()
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: types.TracebackType | None,
    ) -> None:
        """Call ``close()`` on exit."""
        await self.close()

    async def messages(self) -> AsyncGenerator[dict[str, Any]]:
        """Generate messages received from the server."""
        if self.main_loop_task is None:
            raise AssertionError("Main loop task not started")

        # If close() is called during iteration self.main_loop_task can become
        # None, so we need to handle it here
        while self.main_loop_task and not self.main_loop_task.done():
            get_task = asyncio.create_task(self.message_queue.get())

            done, _ = await asyncio.wait(
                (self.main_loop_task, get_task),
                return_when=asyncio.FIRST_COMPLETED,
            )
            if get_task in done:
                payload = get_task.result()
                yield payload
                self.message_queue.task_done()

    async def main_loop(self) -> None:
        """
        Read notifications from the server and add them to ``message_queue``.

        Retries and reconnects are transparently handled.
        """
        reader = tenacity.retry(
            sleep=self._asyncio_sleep,
            wait=tenacity.wait_random(min=1, max=6),
            before=self._tenacity_before_hook,
            retry=tenacity.retry_if_exception_type(
                aiohttp.client_exceptions.ClientError
            )
            | tenacity.retry_if_exception_type(TokenDisabledError)
            | tenacity.retry_if_exception_type(ConnectionError),
        )(self._wait_for_messages)
        await reader()
        logger.debug("Server notification reading loop ended")

    async def _wait_for_messages(self) -> None:
        """Connect to the server and waits for commands."""
        # Enforced by constructor
        assert self.client.token is not None
        headers = {'token': self.client.token}
        async with aiohttp.ClientSession(
            headers={"User-Agent": get_user_agent()}
        ) as client_session:
            self.client_session = client_session

            try:
                async with self.client_session.ws_connect(
                    self.worker_connect_ws_url, headers=headers, heartbeat=60
                ) as ws:
                    logger.info("Websocket connection has been established")
                    self._reset_retry_attempt_number = True
                    msg: aiohttp.WSMessage
                    async for msg in ws:  # pragma: no branch
                        logger.debug(
                            "Incoming websocket message (%s): %s",
                            msg.type.name,
                            str(msg.data),
                        )
                        try:
                            payload = self.client.decode_ws_message(msg)
                        except ServerNotificationDecodeError as exc:
                            logger.warning(
                                "Unhandled websocket notification: %s", exc
                            )
                            continue

                        await self.message_queue.put(payload)

                    logger.info("Server notification channel closed")

                    if not self._shutting_down:
                        # The server disconnected the worker or we lost the
                        # connection
                        logger.warning(
                            "Websocket connection got closed (close_code=%s)",
                            ws.close_code,
                        )
                        self._raise_try_again()

            except TokenDisabledError as exc:
                # The server disconnected the worker because the token is
                # not associated with a worker or is disabled. The server
                # admin should enable the token.
                logger.info(  # noqa: G200
                    'The token is disabled. '
                    'Debusine admin (%s) should enable it. '
                    'Reason: %s. '
                    'Will try again',
                    self.debusine_ws_url,
                    exc,
                )
                raise
            except (
                aiohttp.client_exceptions.ClientConnectorError,
                aiohttp.client_exceptions.WSServerHandshakeError,
            ) as exc:
                logger.info(  # noqa: G200
                    'Error connecting to %s (%s)',
                    self.debusine_ws_url,
                    exc,
                )
                raise
            except ConnectionResetError:
                logger.debug('ConnectionResetError, will reconnect')
                raise

        self.client_session = None
