from asyncio import create_task, Event, sleep
from time import time
from typing import Any, Callable, Union
from logging import Logger, INFO

from aio_pika import connect_robust, Message, DeliveryMode

from aio_pika.abc import AbstractIncomingMessage
from aio_pika.abc import AbstractChannel
from aio_pika.abc import AbstractConnection
from aio_pika.abc import AbstractExchange
from aio_pika.abc import AbstractQueue
from aio_pika.exceptions import AMQPConnectionError

from kernel.rabbitmq.enum.content_type import ContentType
from kernel.rabbitmq.view.authentication_params import AuthenticationParams
from kernel.rabbitmq.view.connection_params import ConnectionParams
from kernel.rabbitmq.exception.connection_timeout import ConnectionTimeout
from kernel.rabbitmq.exception.connection_not_established import ConnectionNotEstablished


class Client:
    __authentication_params: AuthenticationParams
    __connection_params: ConnectionParams
    __logger: Union[None, Logger]
    __connection: Union[None, AbstractConnection] = None
    __stop_event: Union[None, Event] = None

    def __init__(
        self,
        authentication_params: AuthenticationParams,
        connection_params: ConnectionParams,
        logger: Union[None, Logger] = None
    ) -> None:
        self.__authentication_params = authentication_params
        self.__connection_params = connection_params
        self.__logger = logger

    async def connect(self, stop_event: Event, timeout: int = 15) -> None:
        self.__stop_event = stop_event

        ran_at: float = time()

        self.__log(f"Waiting for RabbitMQ connection..")

        while True:
            host = self.__connection_params.host
            port = self.__connection_params.port
            vhost = self.__connection_params.vhost

            try:
                self.__connection = await connect_robust(
                    host=host,
                    port=port,
                    virtualhost=vhost,
                    login=self.__authentication_params.username,
                    password=self.__authentication_params.password,
                )

                self.__log(
                    f"Successfully connected to {host}:{port}{vhost}"
                )

                break
            except AMQPConnectionError:
                seconds_wasted: float = time() - ran_at

                self.__log("Connection failed! Retry after 3s..")

                if seconds_wasted > timeout:
                    raise ConnectionTimeout(
                        f"Connection timeout to {host}:{port}{vhost}"
                    )

                await sleep(3)

        create_task(self.__wait_for_stop())

    async def consume(
        self,
        queue: str,
        exchange: str,
        consumer: Callable[[AbstractIncomingMessage], Any],
        routing_key: Union[None, str] = None,
    ) -> None:
        self.__raise_if_disconnected()
        assert self.__connection is not None

        abstract_channel: AbstractChannel = await self.__connection.channel()
        abstract_exchange: AbstractExchange = await abstract_channel\
            .get_exchange(exchange)
        abstract_queue: AbstractQueue = await abstract_channel\
            .declare_queue(queue, durable=True)

        await abstract_queue.bind(abstract_exchange, routing_key)

        self.__log(
            f"Listening queue {abstract_queue.name} binded to {abstract_exchange.name}..."
        )

        await abstract_queue.consume(consumer)

    async def publish(
        self,
        exchange: str,
        message: str,
        routing_key: str = "",
        content_type: str = ContentType.APPLICATION_JSON
    ) -> None:
        self.__raise_if_disconnected()
        assert self.__connection is not None

        abstract_exchange: AbstractExchange = await self.__connection.channel()\
            .get_exchange(exchange)

        amqp_message = Message(
            body=message.encode(),
            content_type=content_type,
            delivery_mode=DeliveryMode.PERSISTENT,
        )

        await abstract_exchange.publish(amqp_message, routing_key=routing_key)

    def __raise_if_disconnected(self) -> None:
        if self.__connection is None or self.__connection.is_closed:
            raise ConnectionNotEstablished("Connection not established")

    async def __wait_for_stop(self) -> None:
        if self.__stop_event is not None:
            await self.__stop_event.wait()

            self.__log(
                "Received stop signal. Terminating RabbitMQ client.."
            )

            connection_established = self.__connection is not None \
                and not self.__connection.is_closed

            if connection_established:
                await self.__connection.close() # type: ignore

    def __log(self, message: str, level: int = INFO) -> None:
        if self.__logger is not None:
            self.__logger.log(level, message)