import json
import logging
import os
import re
import sys
import uuid

import boto3

BUCKET_NAME = os.getenv('BUCKET_NAME')
OBJECT_PREFIX = os.getenv('OBJECT_PREFIX', '')
OBJECT_FILTER = os.getenv('OBJECT_FILTER', '')

SQS_QUEUE_URL = os.getenv('SQS_QUEUE_URL')
SQS_BATCH_SIZE = os.getenv('SQS_BATCH_SIZE', '10')
SQS_MESSAGE_SIZE = os.getenv('SQS_MESSAGE_SIZE', '10')

CLOUDWATCH_NAMESPACE = os.getenv('CLOUDWATCH_NAMESPACE', 's3restore')

AWS_DEFAULT_REGION = os.environ.get('AWS_DEFAULT_REGION', 'us-east-1')
AWS_REGION = os.environ.get('AWS_REGION', AWS_DEFAULT_REGION)

LOG_LEVEL = os.getenv('LOG_LEVEL', logging.INFO)

s3 = boto3.resource('s3', region_name=AWS_REGION)
sqs = boto3.resource('sqs', region_name=AWS_REGION)
cloudwatch = boto3.resource('cloudwatch', region_name=AWS_REGION)

sqs_queue = None
sqs_batch_buffer = []
sqs_message_size_buffer = []

scanned_object_count = 0
buffered_object_count = 0
queued_object_count = 0


def init():
    logging.basicConfig(
        format='%(asctime)s %(levelname)s %(message)s',
        level=LOG_LEVEL
    )

    if not BUCKET_NAME:
        logging.error('You need to set "BUCKET_NAME" environment variable.')
        sys.exit(1)

    if not SQS_QUEUE_URL:
        logging.error('You need to set "SQS_QUEUE_URL" environment variable.')
        sys.exit(1)

    try:
        global SQS_BATCH_SIZE, SQS_MESSAGE_SIZE
        SQS_BATCH_SIZE = int(SQS_BATCH_SIZE)
        SQS_MESSAGE_SIZE = int(SQS_MESSAGE_SIZE)
    except ValueError:
        logging.error('Incorrect format specified for "SQS_BATCH_SIZE" or '
                      '"SQS_MESSAGE_SIZE" environment variable which should '
                      'be convertible to an integer.')
        sys.exit(1)

    global sqs_queue
    sqs_queue = sqs.Queue(SQS_QUEUE_URL)


def process_bucket():
    global scanned_object_count, buffered_object_count, queued_object_count

    is_truncated = True
    next_key_marker = ''
    next_version_id_marker = ''

    logging.warning('Scanning of bucket "%s" within "%s" object prefix has '
                    'started.', BUCKET_NAME, OBJECT_PREFIX)

    if OBJECT_FILTER:
        logging.warning('Bucket object key filter is set to "%s".',
                        OBJECT_FILTER)
    else:
        logging.warning('Bucket object key filter is not set.')

    while is_truncated:
        if next_key_marker:
            all_objects = s3.meta.client.list_object_versions(
                Bucket=BUCKET_NAME, Prefix=OBJECT_PREFIX,
                KeyMarker=next_key_marker,
                VersionIdMarker=next_version_id_marker)
        else:
            all_objects = s3.meta.client.list_object_versions(
                Bucket=BUCKET_NAME, Prefix=OBJECT_PREFIX)

        if 'Versions' in all_objects:
            all_versions = all_objects['Versions']

            batch_size = len(all_versions)
            if batch_size > 0:
                logging.warning(
                    'Scanning statistics: objects scanned = %d, '
                    'objects queued for analysis = %d, '
                    'objects buffered = %d, current batch size = %d, '
                    'current bucket position = "%s".',
                    scanned_object_count, queued_object_count,
                    buffered_object_count, batch_size, all_versions[0]['Key'])

                publish_metrics()
                scanned_object_count += batch_size

            for version in all_versions:
                if version['IsLatest']:
                    logging.info('Object "%s" is of the latest version, '
                                 'nothing to do.', version['Key'])
                    continue

                if (OBJECT_FILTER
                   and not re.search(OBJECT_FILTER, version['Key'])):
                    logging.info('Object "%s" is filtered out by the object '
                                 'filter, nothing to do.', version['Key'])
                    continue

                process_object(version['Key'])

        is_truncated = all_objects['IsTruncated']
        if is_truncated:
            next_key_marker = all_objects['NextKeyMarker']
            next_version_id_marker = all_objects['NextVersionIdMarker']

    process_queue(True)
    publish_metrics()

    logging.warning('Scanning of bucket "%s" within "%s" object prefix has '
                    'been completed. %d objects have been scanned, %d queued '
                    'for recovery analysis.',
                    BUCKET_NAME, OBJECT_PREFIX,
                    scanned_object_count, queued_object_count)


def process_object(object_key):
    global sqs_message_size_buffer
    global buffered_object_count

    sqs_message_size_buffer.append(object_key)
    buffered_object_count += 1

    logging.info('Object "s3://%s/%s" was queued for recovery analysis.',
                 BUCKET_NAME, object_key)

    process_queue()


def process_queue(flush=False):
    global sqs_queue
    global sqs_batch_buffer, sqs_message_size_buffer
    global buffered_object_count, queued_object_count

    if flush or len(sqs_message_size_buffer) >= SQS_MESSAGE_SIZE:
        content = {
            'bucket_name': BUCKET_NAME,
            'object_keys': sqs_message_size_buffer
        }

        id = str(uuid.uuid4())
        body = json.dumps(content)
        sqs_batch_buffer.append({'Id': id, 'MessageBody': body})

        sqs_message_size_buffer = []

    if flush or len(sqs_batch_buffer) >= SQS_BATCH_SIZE:
        logging.info('Flushing %d objects to SQS.', buffered_object_count)

        sqs_queue.send_messages(Entries=sqs_batch_buffer)

        sqs_batch_buffer = []
        queued_object_count += buffered_object_count
        buffered_object_count = 0


def publish_metrics():
    global scanned_object_count, buffered_object_count, queued_object_count

    dimensions = [
      {
          'Name': 'bucket_name',
          'Value': BUCKET_NAME
      },
      {
          'Name': 'object_prefix',
          'Value': OBJECT_PREFIX if OBJECT_PREFIX else 'None'
      },
      {
          'Name': 'object_filter',
          'Value': OBJECT_FILTER if OBJECT_FILTER else 'None'
      }
    ]

    cloudwatch.meta.client.put_metric_data(
        Namespace=CLOUDWATCH_NAMESPACE,
        MetricData=[
            {
                'MetricName': 'objects-scanned',
                'Dimensions': dimensions,
                'Value': scanned_object_count,
                'Unit': 'Count'
            },
            {
                'MetricName': 'objects-buffered',
                'Dimensions': dimensions,
                'Value': buffered_object_count,
                'Unit': 'Count'
            },
            {
                'MetricName': 'objects-queued-for-analysis',
                'Dimensions': dimensions,
                'Value': queued_object_count,
                'Unit': 'Count'
            }
        ])


if __name__ == '__main__':
    init()
    process_bucket()