import json
import logging
import typing
from collections import Counter
from concurrent.futures import Future

import mock
from google.api_core.exceptions import DeadlineExceeded
from google.cloud import pubsub_v1
from google.cloud.pubsub_v1.proto.pubsub_pb2 import ReceivedMessage

from hedwig.backends.base import HedwigPublisherBaseBackend, HedwigConsumerBaseBackend
from hedwig.backends.import_utils import import_class
from hedwig.conf import settings
from hedwig.models import Message


logger = logging.getLogger(__name__)


# TODO move to dataclasses in py3.7
class GoogleMetadata:
    def __init__(self, ack_id):
        self._ack_id = ack_id

    @property
    def ack_id(self):
        return self._ack_id

    def __eq__(self, o: object) -> bool:
        if not isinstance(o, GoogleMetadata):
            return False
        return self.ack_id == o.ack_id

    def __ne__(self, o: object) -> bool:
        return not self.__eq__(o)

    def __str__(self) -> str:
        return repr(self)

    def __repr__(self) -> str:
        return f'GoogleMetadata(ack_id={self.ack_id})'

    def __hash__(self) -> int:
        return hash((self.ack_id,))


class GooglePubSubAsyncPublisherBackend(HedwigPublisherBaseBackend):
    def __init__(self) -> None:
        self.publisher = pubsub_v1.PublisherClient.from_service_account_file(
            settings.GOOGLE_APPLICATION_CREDENTIALS, batch_settings=settings.HEDWIG_PUBLISHER_GCP_BATCH_SETTINGS
        )

    def publish_to_topic(
        self, topic_path: str, data: bytes, attrs: typing.Optional[typing.Mapping] = None
    ) -> typing.Union[str, Future]:
        """
        Publishes to a Google Pub/Sub topic and returns a future that represents the publish API call. These API calls
        are batched for better performance.

        Note: despite the signature this doesn't return an actual instance of Future class, but an object that conforms
        to Future class. There's no generic type to represent future objects though.
        """
        attrs = attrs or {}
        attrs = dict((str(key), str(value)) for key, value in attrs.items())
        return self.publisher.publish(topic_path, data=data, **attrs)

    def _get_topic_path(self, message: Message) -> str:
        return self.publisher.topic_path(settings.GOOGLE_PUBSUB_PROJECT_ID, f'hedwig-{message.topic}')

    def _mock_queue_message(self, message: Message) -> mock.Mock:
        gcp_message = mock.Mock()
        gcp_message.message = mock.Mock()
        gcp_message.message.data = json.dumps(message.as_dict()).encode('utf8')
        gcp_message.ack_id = 'test-receipt'
        return gcp_message

    def _publish(
        self, message: Message, payload: str, headers: typing.Optional[typing.Mapping] = None
    ) -> typing.Union[str, Future]:
        topic_path = self._get_topic_path(message)
        return self.publish_to_topic(topic_path, payload.encode('utf8'), headers)


class GooglePubSubPublisherBackend(GooglePubSubAsyncPublisherBackend):
    def publish_to_topic(
        self, topic_path: str, data: bytes, attrs: typing.Optional[typing.Mapping] = None
    ) -> typing.Union[str, Future]:
        return typing.cast(Future, super().publish_to_topic(topic_path, data, attrs)).result()


class GooglePubSubConsumerBackend(HedwigConsumerBaseBackend):
    def __init__(self, dlq=False) -> None:
        self.subscriber = pubsub_v1.SubscriberClient.from_service_account_file(settings.GOOGLE_APPLICATION_CREDENTIALS)
        self._publisher = pubsub_v1.PublisherClient.from_service_account_file(settings.GOOGLE_APPLICATION_CREDENTIALS)
        self.message_retry_state: typing.Optional[MessageRetryStateBackend] = None
        self._subscription_path: str = pubsub_v1.SubscriberClient.subscription_path(
            settings.GOOGLE_PUBSUB_PROJECT_ID, f'hedwig-{settings.HEDWIG_QUEUE}{"-dlq" if dlq else ""}'
        )
        self._dlq_topic_path: str = pubsub_v1.PublisherClient.topic_path(
            settings.GOOGLE_PUBSUB_PROJECT_ID, f'hedwig-{settings.HEDWIG_QUEUE}-dlq'
        )
        if settings.HEDWIG_GOOGLE_MESSAGE_RETRY_STATE_BACKEND:
            message_retry_state_cls = import_class(settings.HEDWIG_GOOGLE_MESSAGE_RETRY_STATE_BACKEND)
            self.message_retry_state = message_retry_state_cls()

    def pull_messages(self, num_messages: int = 1, visibility_timeout: int = None) -> typing.List:
        try:
            received_messages = self.subscriber.pull(
                self._subscription_path, num_messages, retry=None, timeout=settings.GOOGLE_PUBSUB_READ_TIMEOUT_S
            ).received_messages
            return received_messages
        except DeadlineExceeded:
            logger.debug(f"Pulling deadline exceeded subscription={self._subscription_path}")
            return []

    def process_message(self, queue_message: ReceivedMessage) -> None:
        try:
            self.message_handler(queue_message.message.data.decode(), GoogleMetadata(queue_message.ack_id))
        except Exception:
            if self._can_reprocess_message(queue_message):
                raise

    def delete_message(self, queue_message: ReceivedMessage) -> None:
        self.subscriber.acknowledge(self._subscription_path, [queue_message.ack_id])

    @staticmethod
    def pre_process_hook_kwargs(queue_message: ReceivedMessage) -> dict:
        return {'google_pubsub_message': queue_message}

    @staticmethod
    def post_process_hook_kwargs(queue_message: ReceivedMessage) -> dict:
        return {'google_pubsub_message': queue_message}

    def extend_visibility_timeout(self, visibility_timeout_s: int, metadata: GoogleMetadata) -> None:
        """
        Extends visibility timeout of a message on a given priority queue for long running tasks.
        """
        if visibility_timeout_s < 0 or visibility_timeout_s > 600:
            raise ValueError("Invalid visibility_timeout_s")
        self.subscriber.modify_ack_deadline(self._subscription_path, [metadata.ack_id], visibility_timeout_s)

    def requeue_dead_letter(self, num_messages: int = 10, visibility_timeout: int = None) -> None:
        """
        Re-queues everything in the Hedwig DLQ back into the Hedwig queue.

        :param num_messages: Maximum number of messages to fetch in one call. Defaults to 10.
        :param visibility_timeout: The number of seconds the message should remain invisible to other queue readers.
        Defaults to None, which is queue default
        """
        topic_path = pubsub_v1.PublisherClient.topic_path(
            settings.GOOGLE_PUBSUB_PROJECT_ID, f'hedwig-{settings.HEDWIG_QUEUE}'
        )

        logging.info("Re-queueing messages from {} to {}".format(self._subscription_path, topic_path))
        while True:
            queue_messages = self.pull_messages(num_messages=num_messages, visibility_timeout=visibility_timeout)
            if not queue_messages:
                break

            logging.info("got {} messages from dlq".format(len(queue_messages)))
            for queue_message in queue_messages:
                try:
                    if visibility_timeout:
                        self.subscriber.modify_ack_deadline([queue_message.ack_id], visibility_timeout)

                    self._publisher.publish(topic_path, data=queue_message.message.data, **queue_message.attributes)
                    logger.debug(
                        'Re-queued message from DLQ {} to {}'.format(self._subscription_path, topic_path),
                        extra={'message_id': queue_message.message_id},
                    )

                    self.delete_message(queue_message)
                except Exception:
                    logger.exception(
                        'Exception in requeue message from {} to {}'.format(self._subscription_path, topic_path)
                    )

            logging.info("Re-queued {} messages".format(len(queue_messages)))

    def _can_reprocess_message(self, queue_message: ReceivedMessage) -> bool:
        if not self.message_retry_state:
            return True

        try:
            self.message_retry_state.inc(queue_message.message.message_id, self._subscription_path)
            return True
        except MaxRetriesExceededError:
            self._move_message_to_dlq(queue_message)
        return False

    def _move_message_to_dlq(self, queue_message: ReceivedMessage) -> None:
        self._publisher.publish(self._dlq_topic_path, queue_message.message.data, **queue_message.attributes)
        logger.debug('Sent message to DLQ', extra={'message_id': queue_message.message_id})


class MaxRetriesExceededError(Exception):
    pass


class MessageRetryStateBackend:
    def __init__(self) -> None:
        self.max_tries = settings.HEDWIG_GOOGLE_MESSAGE_MAX_RETRIES

    def inc(self, message_id: str, queue_name: str) -> None:
        raise NotImplementedError

    @staticmethod
    def _get_hash(message_id: str, queue_name: str) -> str:
        return f"{queue_name}-{message_id}"


class MessageRetryStateLocMem(MessageRetryStateBackend):
    DB: typing.Counter = Counter()

    def inc(self, message_id: str, queue_name: str) -> None:
        key = self._get_hash(message_id, queue_name)
        self.DB[key] += 1
        if self.DB[key] >= self.max_tries:
            raise MaxRetriesExceededError


class MessageRetryStateRedis(MessageRetryStateBackend):
    def __init__(self) -> None:
        import redis

        super().__init__()
        self.client = redis.from_url(settings.HEDWIG_GOOGLE_MESSAGE_RETRY_STATE_REDIS_URL)

    def inc(self, message_id: str, queue_name: str) -> None:
        key = self._get_hash(message_id, queue_name)
        value = self.client.incr(key)
        if value >= self.max_tries:
            self.client.expire(key, 0)
            raise MaxRetriesExceededError
