import json
import logging
from functools import partial
from typing import Optional, Union, Tuple

import pika
from django.conf import settings
from pika.adapters.blocking_connection import BlockingChannel
from pika.exchange_type import ExchangeType
from rest_framework import status
from rest_framework.exceptions import APIException

from devices.models.sensor import ArmaSensor
from devices.services.sensor.enums import SystemMessage, ZeekMessage, VectorMessage, SuricataMessage

_log = logging.getLogger(__name__)


class SensorResponseException(APIException):
    status_code = status.HTTP_400_BAD_REQUEST


class SensorManagement:
    """Service for sending message and receiving response using RabbitMQ"""

    # Префикс, с которого начинаются все очереди для сенсоров. Каждый сенсор слушает свою очередь
    SENSOR_QUEUE_PREFIX = 'sensor_'

    # "Обменник", к которому можно подключить очередь сенсора.
    # Консоль может отправить одно сообщение в этот "обменник" и оно будет доставлено во все подключённые очереди
    SENSORS_EXCHANGE = 'sensors'

    # если отправляется сообщение, на которое ожидается ответ, то в этом атрибуте хранится список таких ответов
    _rpc_responses = tuple()

    def __init__(self, proceed=True):

        self.connection, self.channel = self._get_connection_data()

        # Флаг означает, что ожидание и обработка ответов от сенсора ещё продолжается
        self._proceed = proceed

    def send_message(self,
                     sensor: ArmaSensor = None,
                     message_type: Union[
                         SystemMessage, VectorMessage, ZeekMessage, SuricataMessage] = SystemMessage.ping,
                     body: Union[dict, bytes, str] = None,
                     wait_response: bool = False,
                     decode_response: bool = True,
                     time_limit: int = 5) -> Optional[dict]:
        """Method for sending messages.

        If we requested a response from the sensor, then it returns  dict(). Otherwise, it will return nothing.
        Throws an exception if the response from the sensor cannot be decoded.

        :param sensor: ArmaSensor object. Specify if you want to send the message only to its queue.
            If it is not transmitted, the message will be sent to all sensors.
        :param message_type: A string that specifies what type of message will be sent.
            Depending on this type, different logic will be called on the sensor side.
            If the sensor does not know how to process any type, then it will return an error
        :param body: The payload that can be sent along with the message.
            Used on the sensor side. We only encode and send
        :param wait_response: A flag that determines whether we are waiting for a response to the message or not.
            If False, then we send the message and immediately close the connection.
            Otherwise, we block the process until we wait for a response.
        :param decode_response: Whether to decode the message that is returned from the sensor
        :param time_limit:
            If we send a message and a response is required, how long to wait for it. Specified in seconds.
            If we sent a message to all sensors, but after a while not all of them answered,
            then we will return only the received response
        :return:
        """

        # Шаг 1. Определяем куда будем отправлять сообщение. Если передан сенсор, то только ему. Иначе отправим всем
        if sensor:
            queue = f'{self.SENSOR_QUEUE_PREFIX}{sensor.uuid}'  # очередь сенсора, которую будет слушать только он
            exchange = ''
            required_responses = 1  # если отправляем сообщение только одному сенсору, то мы ожидаем только один ответ
        else:
            exchange, queue = self._get_sensors_exchange(), ''
            required_responses = ArmaSensor.objects.count()  # ожидаем количество ответов, по количеству сенсоров

        # Шаг 2. Готовим сообщение к отправке
        # 2.1 Приводим необязательную полезную нагрузку к байтам
        bytes_body = self._convert_to_bytes(body)

        # 2.2 Готовим мета-информацию (свойства) для сообщения
        properties = {
            # Тип сообщения. по нему сенсор определит что именно ему нужно сделать
            'type': message_type.value,
            # Как долго отправленное сообщение будет храниться в очереди.
            # По истечение времени если его никто не получит, то оно будет удалено
            'expiration': '5000'
        }

        # 2.3 Если мы отправляем сообщение и хотим получить на него ответ, то должны указать:
        # - в какую очередь сенсор должен направить ответ
        # - какая функция будет вызываться при получении ответа в эту очередь
        if wait_response:
            on_message_callback = partial(self._response_callback, decode_response=decode_response)
            properties['reply_to'] = self._prepare_callback_queue(on_message_callback)

        # конвертируем подготовленные свойства сообщения в особый формат
        properties = pika.BasicProperties(**properties)

        # Шаг 3. Отправляем сообщение. Указываем получателей, тело сообщения и его свойства
        self.channel.basic_publish(exchange=exchange, routing_key=queue, body=bytes_body, properties=properties)

        # Шаг 4. Получаем ответы от сенсора. Если мы их не запрашивали, то из функции вернётся None
        response = self._get_responses(sensor, wait_response, required_responses, time_limit)

        # Шаг 5. Закрываем соединение
        self.connection.close()

        return response

    def _get_responses(self, sensor, wait_response, required_responses, time_limit) -> Union[list, dict]:
        """A function that counts messages and converts them to json format"""

        def _stop_proceed_after_timeout():
            self._proceed = False

        # Определяем функцию и через какое время она будет вызвана.
        # Функция выключит флаг обработки и условие 'while' ниже не будет выполняться.
        # Этот метод не блокирует выполнение
        self.connection.call_later(time_limit, _stop_proceed_after_timeout)

        # псевдокод:
        # пока мы ожидаем ответа и получены не все требуемые ответы и время ещё не вышло: ...
        while wait_response and (len(self._rpc_responses) < required_responses) and self._proceed:
            self.connection.process_data_events()

        # когда будет произведён выход из цикла (после получения всех ответов или по истечение времени) мы начнём
        # обработку ответов
        if wait_response and sensor:
            # если мы ожидаем ответ только от одного сенсора и этот ответ получен,
            # то мы проверяем его статус и возвращаем json ответ
            if self._rpc_responses:
                response = self._check_response_status(self._rpc_responses[0])
                return response
            # если ответ не получен, то возбуждаем исключение
            else:
                raise SensorResponseException({
                    'status': 'error',
                    'detail': 'sensor did not send a response'
                })
        elif wait_response:
            # Если ожидаем ответ от всех сенсоров и закончилось время обработки, то возвращаем всё, что успели получить.
            return self._rpc_responses

    def _check_response_status(self, response):
        """Check response status"""

        # Проверяем статус ответа. При необходимости возбуждаем исключение, чтобы на фронт был отдан ответ с 400 кодом
        # Данное поле считать соглашением между консолью и сенсором
        # Если сенсор не пришлёт это поле или пришлёт что-то другое, то считается, что это корректный ответ без ошибки
        if isinstance(response, dict) and response.get('status', None) == 'error':
            raise SensorResponseException(response)
        return response

    def _get_sensors_exchange(self) -> str:
        """Determine the exchanger and its type."""

        # 'fanout' означает, что сообщение будет отправлено во все очереди, связанные с обменником
        self.channel.exchange_declare(self.SENSORS_EXCHANGE, exchange_type=ExchangeType.fanout.value, durable=True)
        return self.SENSORS_EXCHANGE

    def _prepare_callback_queue(self, callback_func) -> str:
        """We create a queue in which we expect a response"""

        # Создаём одноразовую очередь для ожидания в ней ответа. Она удалится после закрытия соединения
        result = self.channel.queue_declare(queue='', exclusive=True)
        callback_queue = result.method.queue

        # Указываем функцию, которая будет обрабатывать сообщения во временной очереди
        self.channel.basic_consume(queue=callback_queue, on_message_callback=callback_func, auto_ack=True)

        return callback_queue

    @staticmethod
    def _convert_to_bytes(body) -> bytes:
        """Convert body to bytes"""

        if body is None:
            body = b''
        if isinstance(body, bytes):
            return body
        try:
            str_json = json.dumps(body)
            bytes_body = str_json.encode()
            return bytes_body
        except TypeError:
            raise SensorResponseException({
                'status': 'error',
                'detail': 'unable to encode data to send'
            })

    def _response_callback(self, channel, method, properties, body, decode_response: bool):
        """A function that processes messages in the response queue

        Receives response, converts to json and writes to attribute self._rpc_responses
        """

        response = self._prepare_response(body, decode_response)
        self._rpc_responses += (response, )

    def _prepare_response(self, rpc_response: Optional[bytes], decode_rcp_response: bool) -> Union[dict, bytes]:

        if not rpc_response:
            response = {
                'status': 'error',
                'detail': 'sensor doesnt return response'
            }
            return response if decode_rcp_response else json.dumps(response)

        if decode_rcp_response:
            response = self._decode_rpc_response(rpc_response)
        else:
            response = rpc_response
        return response

    @staticmethod
    def _decode_rpc_response(rcp_response: bytes) -> dict:
        try:
            response = json.loads(rcp_response)
        except json.JSONDecodeError:
            response = {
                'status': 'error',
                'detail': 'failed to decode sensor response to json'
            }
        return response

    @staticmethod
    def _get_connection_data() -> Tuple[pika.BlockingConnection, BlockingChannel]:
        """Create connection to RabbitMQ"""

        connection = pika.BlockingConnection(pika.ConnectionParameters(host=settings.RABBIT_HOST,
                                                                       port=settings.RABBIT_PORT))
        channel = connection.channel()

        return connection, channel
