"""Client for adding ingest requests to the ingest queue."""

from __future__ import absolute_import

import base64
import itertools
import json
import logging

import boto3

from . import errors
from .proto import ingest_requests_pb2


logger = logging.getLogger(__name__)

MAX_MESSAGE_BYTES = 262144
MAX_BATCH_MESSAGES = 10
MAX_BATCH_BYTES = 262144


class IngestQueue(object):
    """AWS SQS Queue used to ingest some kinds of data into RPS.

    Args:
      queue_name: string name of SQS queue_name
      region: string name of the AWS region where the queue is

    Properties:
      queue_name: same as constructor arg
      sqs: boto3 SQS resource
      queue: boto3 SQS Queue resource
    """

    def __init__(self, queue_name, region='us-west-2'):
        self.sqs = boto3.resource('sqs', region_name=region)
        self.queue_name = queue_name
        self._queue = None

    @property
    def queue(self):
        if not self._queue:
            self._queue = self.sqs.get_queue_by_name(QueueName=self.queue_name)
        return self._queue

    def AddRequests(self, requests):
        """Marshal requests and send them to the ingest queue.

        The requests will be batched up.

        Args:
          requests: a iterable of IngestRequest protobuf objects.

        Raises:
          IngestQueueAddRequestsFailures on failure, which will have a
          "failures" attribute with a list of BatchResultErrorEntry items.
        """
        entries = (self._MakeEntry(i, r) for i, r in enumerate(requests))

        for batch in self._Batcher(entries):
            failures = self._SendBatch(batch)
            for failed in failures:
                logger.error('Failed to enqueue request: %s',
                             json.dumps(failed))
            if failures:
                e = errors.IngestQueueAddRequestsFailures()
                e.failures = failures
                raise e

    def _MakeEntry(self, id, request):
        """Returns json-like dict used as args for SQS send_message."""
        message_body = base64.b64encode(request.SerializeToString())
        if len(message_body) > MAX_MESSAGE_BYTES:
            raise errors.IngestQueueRequestTooLarge()
        return {
          'Id': str(id),
          'MessageBody': message_body,
        }

    def _Batcher(self, entries):
        class Batcher(object):

            def __init__(self):
                self.messages = 0
                self.bytes = 0
                self.key = 0

            def __call__(self, x):
                if (self.messages + 1 > MAX_BATCH_MESSAGES or
                        self.bytes + len(x['MessageBody']) > MAX_BATCH_BYTES):
                    self.messages = 0
                    self.bytes = 0
                    self.key += 1
                self.messages += 1
                self.bytes += len(x['MessageBody'])
                return self.key

        for _, batch in itertools.groupby(iter(entries), Batcher()):
            yield batch

    def _SendBatch(self, entries):
        response = self.queue.send_messages(Entries=list(entries))
        successes = response.get('Successful', [])
        failures = response.get('Failed', [])
        logger.info('Enqueued %d successfully, %d failed',
                     len(successes), len(failures))
        return failures


def IngestBlueprintRequest(github_repository, blueprint_path,
                           github_host=None):
    """Constructs a IngestRequest for ingesting blueprints.

    Args:
      github_repository: string like "dta/rockpaperscissors"
      blueprint_path: path of blueprint file, relative to root of the repo
      github_host: optional, host of the github repository

    Returns:
      IngestRequest protobuf object
    """
    request = ingest_requests_pb2.IngestRequest()
    ingest_request = request.ingest_blueprint_request
    ingest_request.source.repository.github_repository.name = github_repository
    if github_host:
        ingest_request.source.repository.github_repository.host = github_host
    ingest_request.source.path = blueprint_path
    return request


def IngestGitHubStatsRequest(github_repository, commit_sha, github_host=None):
    """Constructs a IngestRequest for ingesting GitHub commit stats.

    Args:
      github_repository: string like "dta/rockpaperscissors"
      commit_sha: string commit id/sha to collect stats for
      github_host: optional, host of the github repository

    Returns:
      IngestRequest protobuf object
    """
    request = ingest_requests_pb2.IngestRequest()
    ingest_request = request.ingest_github_stats_request
    ingest_request.github_repository.name = github_repository
    if github_host:
        ingest_request.github_repository.host = github_host
    ingest_request.commit_sha = commit_sha
    return request
