"""Script to move existing statement attachments to per label directory.

Please read README.
"""

import argparse
import datetime
import logging
import sys
import os

from boto3.dynamodb import conditions
from botocore import exceptions

import common_config
import dynamodb
import s3
import util

LOGGER_INFO_FOMRAT = '%(asctime)s %(levelno)s: %(message)s'
LOGGER_ERROR_FORMAT = '%(asctime)s %(levelno)s %(funcName)s: %(message)s'
TABLE_NAME = '{env}_statement_attachments_success'
ATTRS_TO_GET = (
    'label_type_id_period_id,'
    'file_name,'
    'file_key,'
    'bucket_name'
)
EXECUTION_FINISHED_MESSAGE = (
    'Execution finished. Scanned: {scanned_count} item(s), found: {count} \n'
    'Successfully moved: {successful_count} item(s)\n'
    'Had troubles with {errors_count} item(s)\n'
    'Errors: {errors}'
)


def is_month_item(item):
    """Check if provided item is for month accounting period.

    Args:
        item (dict): DynamoDB scan results, attachment details

    Returns:
        bool: True if item corresponds to month record in DynamoDB else False

    """
    key = item['label_type_id_period_id']
    return len(key.split('_')) == 2


def get_dest_file_key(label_type_id, file_key, dest_dir):
    """Build S3 destination key from file_key and dest_dir.

    Args:
        label_type_id (str): label type and id e.g. L4242
        file_key (str): original flie key
        dest_dir (str): destination directory to put file

    Returns:
        str: S3 destination file key

    """
    full_file_name = util.get_file_name_from_key(file_key)
    label_dir = os.path.join(dest_dir, label_type_id)
    return os.path.join(label_dir, full_file_name)


def get_approximate_table_count(env):
    """Get approximate item count for DynamoDB table.

    Args:
        env (str): environment

    Returns:
        int: item count, might reflect ~6 hrs old data

    """
    table_name = TABLE_NAME.format(env=env)
    table = dynamodb.get_dynamodb_table(table_name)
    return table.item_count


def get_estimated_time(item_count, limit):
    """Calculate estimated execution time.

    Args:
        item_count (int): DynamoDB table item count
        limit (int): DynamoDB scan operation limit

    Returns:
        str: estimated execution time

    """
    if item_count == 0 or limit == 0:
        seconds = 0
    else:
        seconds = item_count / limit * 2
    return str(datetime.timedelta(seconds=seconds))


def build_update_kwargs(dynamodb_key, s3_key):
    """Build update kwargs to be used in dynamodb.update_item function.

    Args:
        dynamodb_key (dict): DynamoDB primary and sort keys
        s3_key (str): new S3 key to store

    Returns:
        dict: args suitable for dynamodb.update_existing_item function

    """
    update_data = {
        'Key': dynamodb_key,
        'UpdateExpression': 'SET file_key = :file_key',
        'ExpressionAttributeValues': {':file_key': s3_key},
    }
    return update_data


def get_primary_keys(key_month, file_name):
    """Build primary keys suitable for DynamoDB update operation.

    Args:
        key_month (str): DynamoDB key for month record
        file_name (str): file name (without label id and period)

    Returns:
        list: list of dict objects with primary keys

    """
    label_key, period_id = key_month.split('_')
    quarter_periods = util.get_quarter_periods_by_period(int(period_id))
    key_quarter = util.get_attachment_primary_key_value(
        label_key[0], label_key[1:], quarter_periods)
    update_keys = [
        {
            'label_type_id_period_id': key_month,
            'file_name': file_name},
        {
            'label_type_id_period_id': key_quarter,
            'file_name': util.format_quarter_file_name(file_name, period_id)}
    ]
    return update_keys


def move_attachment_and_update_record(
        env, bucket, source_key, dest_key, dynamodb_keys):
    """Move S3 file and update the DynamoDB records for the attachment.

    Args:
        env (str): environment
        bucket (str): bucket name
        source_key (str): S3 source key
        dest_key (str): S3 destination key
        dynamodb_keys (list): list of dicts DynamoDB keys to update

    Returns:
        bool: move and update operations result

    """
    copy_successful = s3.copy_object(
        bucket=bucket,
        source_key=source_key,
        dest_key=dest_key)
    if not copy_successful:
        return False
    table_name = TABLE_NAME.format(env=env)
    table = dynamodb.get_dynamodb_table(table_name)
    update_successful = True
    for dynamodb_key in dynamodb_keys:
        update_data = build_update_kwargs(dynamodb_key, dest_key)
        try:
            dynamodb.update_existing_item(table, update_data)
        except (exceptions.ClientError, exceptions.BotoCoreError) as e:
            update_successful = False
            common_config.logger.exception(
                'Error updating key: {}'.format(dynamodb_key))
    if update_successful:
        delete_successful = s3.delete_object(bucket, source_key)
        if not delete_successful:
            common_config.logger.info('Could not delete {}'.format(source_key))
    return True


def process_items(env, items, dest_dir):
    """Move S3 files to destination dir and update the DynamoDB records.

    Args:
        env (str): environment
        items (list): DynamoDB scan results, statement attachment records
        dest_dir (str): destination directory to move S3 files

    Returns:
        dict: processing result {'successful_count': int, 'errors: [dict]}

    """
    result = {
        'successful_count': 0,
        'errors': []
    }
    for item in items:
        # skip quarter items, move only month records, update both
        if not is_month_item(item):
            continue
        key_month = item['label_type_id_period_id']
        bucket = item['bucket_name']
        source_key = item['file_key']
        label_type_id = key_month.split('_')[0]

        keys_to_update = get_primary_keys(key_month, item['file_name'])
        dest_key = get_dest_file_key(label_type_id, source_key, dest_dir)

        move_successful = move_attachment_and_update_record(
            env=env,
            bucket=bucket,
            source_key=source_key,
            dest_key=dest_key,
            dynamodb_keys=keys_to_update)
        if move_successful:
            result['successful_count'] += 1
        else:
            result['errors'].append(item)
    return result


def get_data_from_dynamodb(env, source, limit):
    """Retrieve the data from DynamoDB table.

    Args:
        env (str): environment
        source (str): file_key prefix
        limit (int): scan operation limit

    Returns:
        generator: Scan operation response

    """
    table_name = TABLE_NAME.format(env=env)
    table = dynamodb.get_dynamodb_table(table_name)
    condition = conditions.Attr('file_key').begins_with(source)
    scan_kwargs = {
        'Select': 'SPECIFIC_ATTRIBUTES',
        'ProjectionExpression': ATTRS_TO_GET,
        'FilterExpression': condition,
        'ConsistentRead': True,
        'Limit': limit,
        'ReturnConsumedCapacity': 'TOTAL'
    }
    for resp in dynamodb.get_scan_results(table, scan_kwargs):
        common_config.logger.info('Consumed capacity: {}'.format(
            resp.get('ConsumedCapacity')))
        yield resp


def move_valid_statement_attachments(env, source_dir, dest_dir, limit):
    """Move statement from source directory to destination directory."""
    result = {
        'Count': 0,
        'ScannedCount': 0,
        'successful_count': 0,
        'errors': []
    }
    for chunk in get_data_from_dynamodb(env, source_dir, limit):
        if 'Items' in chunk:
            process_result = process_items(
                env=env, items=chunk['Items'], dest_dir=dest_dir)
            result['successful_count'] += process_result['successful_count']
            result['errors'].extend(process_result['errors'])
        result['Count'] += chunk['Count']
        result['ScannedCount'] += chunk['ScannedCount']
    return result


def configure_logger(logger):
    """Configure existing logger.

    Args:
        logger (logging.logger): logger instance

    """
    # info
    info_handler = logging.StreamHandler(stream=sys.stdout)
    info_formatter = logging.Formatter(LOGGER_INFO_FOMRAT)
    info_handler.setFormatter(info_formatter)
    info_handler.setLevel(logging.INFO)
    logger.addHandler(info_handler)
    # error
    error_handler = logging.StreamHandler(stream=sys.stderr)
    error_formatter = logging.Formatter(LOGGER_ERROR_FORMAT)
    error_handler.setFormatter(error_formatter)
    error_handler.setLevel(logging.ERROR)
    logger.addHandler(error_handler)


def main(args):
    """Entry point of the script.

    Args:
        args (): argparse arguments

    """
    # setup logging
    common_config.logger.setLevel(args.log_level.upper())
    configure_logger(common_config.logger)

    dest_dir = args.dest_dir
    if not dest_dir:
        dest_dir = '{env}-per-label-statement-attachments'.format(args.env)

    item_count = get_approximate_table_count(args.env)
    limit = args.scan_limit
    esitmated_time = get_estimated_time(item_count, limit)
    common_config.logger.info('Estimated execution time: {}'.format(
        esitmated_time))

    # start processing
    move_results = move_valid_statement_attachments(
        env=args.env,
        source_dir=args.source_dir,
        dest_dir=dest_dir,
        limit=limit)
    # log results
    errors_msg = '\n'.join(str(error) for error in move_results['errors'])
    msg = EXECUTION_FINISHED_MESSAGE.format(
        scanned_count=move_results['ScannedCount'],
        count=move_results['Count'],
        successful_count=move_results['successful_count'],
        errors_count=len(move_results['errors']),
        errors=errors_msg or 'No errors.'
    )
    common_config.logger.info(msg)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Move statement attachments from source dir to dest dir.')
    parser.add_argument(
        '-e', '--env', default='dev',
        help='Environment, one of: dev, qa, prod',
        choices=['dev', 'qa', 'prod'])
    parser.add_argument(
        '-s', '--source_dir', default='',
        help='Source directory (file key prefix).')
    parser.add_argument(
        '-d', '--dest_dir', default='',
        help=(
            'Destination directory. '
            'Defaults to `env`-per-label-statement-attachments'))
    parser.add_argument(
        '--scan_limit', type=int, default=100,
        help=(
            'DynamoDB scan operation limit. '
            'Defaults to 100 records per request.'))
    parser.add_argument(
        '-l', '--log_level',
        help='logging level', default='error',
        choices=['critical', 'error', 'warning', 'info', 'debug'])
    args = parser.parse_args()
    main(args)