"""Kickoff re-compilations and deliveries based on S3 file."""

import argparse
import json
import os
import time

import boto3

import botocore
from botocore.exceptions import ClientError
from botocore.exceptions import ConnectionClosedError

import requests


# setup aws account configuration
ENVIRONMENT = os.environ['ENVIRONMENT']
if ENVIRONMENT == 'dev':
    AWS_ACCOUNT = '103233932089'
elif ENVIRONMENT in ('prod', 'qa'):
    AWS_ACCOUNT = '437795906767'
else:
    raise Exception('unknown environment {ENVIRONMENT}')

# verify credentials are the same as ENVIRONMENT expects
CREDS_ACCOUNT = boto3.client('sts').get_caller_identity().get('Account')
if CREDS_ACCOUNT != AWS_ACCOUNT:
    print(f'{ENVIRONMENT} != {CREDS_ACCOUNT}')
    raise Exception('Active AWS credentials are for a different account!')

# setup function name/arn
FB_DELIVERY_SFN_ARN = f'arn:aws:states:us-east-1:{AWS_ACCOUNT}:stateMachine:{ENVIRONMENT}-sr-delivery-sfn'  # noqa:E501
ADD_VERSION_LAMBDA_NAME = f'{ENVIRONMENT}-lambda-sr-add-version'

# pick s3 bucket by environment
S3_CONFIG = {
    'dev': {
        'bucket': 'test-orcd-bucket',
        'prefix': 'dev-interm-sr-delivery/'
    },
    'qa': {
        'bucket': 'qa-sr-interm-delivery-bucket',
        'prefix': ''
    },
    'prod': {
        'bucket': 'prod-sr-interm-delivery-bucket',
        'prefix': ''
    }
}
BUCKET = S3_CONFIG[ENVIRONMENT]['bucket']
PREFIX = S3_CONFIG[ENVIRONMENT]['prefix']

# static sub direction names
SUB_DIR_READY = 'ready'
SUB_DIR_PROCESSING = 'processing'
SUB_DIR_COMPLETE = 'complete'

# boto clients
s3_client = boto3.client('s3')
lambda_client = boto3.client(
    'lambda',
    config=botocore.config.Config(
        retries={'max_attempts': 0},
        read_timeout=900  # wait for max lambda runtime
    )
)
sfn_client = boto3.client('stepfunctions')


class RetryableException(Exception):
    """Retryable operation."""

    def __init__(self, message):
        """Init."""
        super().__init__(message)


def main():
    """Entrypoint."""
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        '--dir',
        type=str,
        choices=[SUB_DIR_READY, SUB_DIR_PROCESSING, SUB_DIR_COMPLETE],
        default=SUB_DIR_READY,
        help='Sub directory to look for files in'
    )
    parser.add_argument(
        '--filename',
        type=str,
        help='Specific file to process, default to oldest in sub dir'
    )
    parser.add_argument(
        '--since',
        type=int,
        help='UPC to start processing from'
    )
    parser.add_argument(
        '--throttle',
        type=int,
        default=10,
        help='Max running delivery SFNs'
    )
    parser.add_argument(
        '--backoff',
        type=int,
        default=180,
        help='Time in seconds to sleep on error'
    )
    parser.add_argument(
        '--retries',
        type=int,
        default=5,
        help='Max number of times to retry errors'
    )
    parser.add_argument(
        '-f',
        action='store_true',
        help='Force compilation of sound recodings'
    )
    args = parser.parse_args()
    print(args)

    sub_dir = args.dir
    arg_filename = args.filename
    since_upc = args.since
    throttle = args.throttle
    backoff = args.backoff
    max_retries = args.retries
    force_compile = args.f

    # read files from sub directory
    files = list_files(sub_dir)
    print(f'Found {len(files)} file(s) to process')
    if not files:
        exit()

    # select file to process
    filenames = [
        x['Key'].split('/')[-1]
        for x in sorted(files, key=lambda d: d['LastModified'])
    ]
    filename = filenames[0]
    if arg_filename:
        if arg_filename not in filenames:
            print(f'Unable to find {arg_filename} in {sub_dir}')
            exit(1)
        filename = arg_filename
    print(f'Picked {filename} to load and process')

    # move file to processing sub dir, if not there already
    if sub_dir != SUB_DIR_PROCESSING:
        move_file(filename, sub_dir, SUB_DIR_PROCESSING)

    # load file into memory, split by newline, validate
    data = load_file(filename, SUB_DIR_PROCESSING).replace('\r', '')
    try:
        upcs = [
            int(x)
            for x in data.split('\n')
            if x
        ]
    except ValueError as e:
        print(e)
        exit(1)

    # start from specific UPC if param set
    if since_upc:
        if since_upc not in upcs:
            print(f'Unable to find {since_upc} in {filename}')
            exit(1)
        upcs = upcs[upcs.index(since_upc):]

    # process each line
    num_processed = 0
    total_upcs = len(upcs)
    print(f'Found {total_upcs} UPC(s) to process')
    for upc in upcs:
        retries = max_retries
        while True:
            if retries == 0:
                print(f'Retried {max_retries} times.. Stopping.')
                break
            try:
                process_upc(upc, throttle, force_compile)
                num_processed += 1
                percent_processed = (num_processed / total_upcs) * 100
                print(f'Processed {num_processed} / {total_upcs} : {percent_processed}%')  # noqa:E501
                break
            except RetryableException as re:
                print(f'Retryable Exception.. {re}')
                retries -= 1
                time.sleep(backoff)
            except Exception as e:
                print(f'Non-Retryable Exception.. {e}')
                raise e

    # move file to completed sub dir
    move_file(filename, SUB_DIR_PROCESSING, SUB_DIR_COMPLETE)

    # all done
    print('Complete!')
    exit()


def get_sr_data(upc):
    """Retrieve OSR data based on UPC."""
    print('Retrieving OSR data...')
    params = {'upcs': upc}
    return requests.get(f'https://{ENVIRONMENT}-ows-sound-recordings.theorchard.io/sound_recordings', params=params)  # noqa:E501


def get_sr_version_data(osrs, force_compile):
    """Retrieve latest OSR versions."""
    print('Retrieving OSRs versions data...')
    data = []
    for sr in osrs:
        osr_id = sr['id']
        osr_results = []

        # try to fetch latest osr data
        if not force_compile:
            response = requests.get(f'https://{ENVIRONMENT}-ows-sound-recordings.theorchard.io/sound_recordings/{osr_id}/versions/latest')  # noqa:E501
            if response.status_code != 200:
                if response.status_code in (500, 502, 503, 504):
                    raise RetryableException(f'Error retrieving version for OSR ID: {osr_id}')  # noqa:E501
                elif response.status_code == 404:
                    print(f'Missing compiled version for OSR ID: {osr_id}')
                else:
                    raise Exception(f'Failed to retrieve version for OSR ID: {osr_id}. HTTP code: {response.status_code}')  # noqa:E501
            else:
                print(f'Found compiled version for OSR ID: {osr_id}')
                osr_results.append((osr_id, response.json()['version_id']))

        # recompile latest osr data if not exists or forced
        if force_compile or not osr_results:
            compiled_results = invoke_lambda_compilation(osr_id)
            for result in compiled_results:
                osr_results.append((osr_id, result['version_id']))

        # add osr results to list
        data += [
            {
                'sound_recording_id': x[0],
                'version_id': x[1]
            }
            for x in osr_results
        ]

    return data


def process_upc(upc, throttle, force_compile):
    """Call functions to process UPC."""
    print(f'Processing {upc}...')

    ows_sr_data_response = get_sr_data(upc)
    if ows_sr_data_response.status_code != 200:
        if ows_sr_data_response.status_code in (500, 502, 503, 504):
            raise RetryableException(f'Error retrieving OSR for UPC: {upc}')
        raise Exception(f'Unexpected HTTP code {ows_sr_data_response.status_code}')  # noqa:E501

    # no backfill or v2 assets found (product needs a migration)
    sr_data = ows_sr_data_response.json()
    if not sr_data:
        print(f'OSR(s) not found for upc: {upc}')

        # kickoff execution that will fail, to create log of this event
        err_exc_name = f'interm-error-{upc}'
        err_payload = {
            'sound_recording': {
                'id': None,
                'version': None
            },
            'upload_asset': False
        }
        start_sfn_execution(err_exc_name, err_payload, throttle)
        return

    ows_sr_version_response = get_sr_version_data(sr_data, force_compile)
    if not ows_sr_version_response:
        raise Exception('OSR(s) versions not found')

    upload = True  # TODO - always upload?
    for sr in ows_sr_version_response:
        v_id = sr['version_id']
        sr_id = sr['sound_recording_id']
        exc_name = f'interm-{sr_id}'
        payload = {
            'sound_recording': {
                'id': sr_id,
                'version': v_id
            },
            'upload_asset': upload
        }
        start_sfn_execution(exc_name, payload, throttle)


def start_sfn_execution(exc_name, payload, throttle):
    """Start new sfn execution, throttling if too many running."""
    while True:
        sfn_list_response = sfn_client.list_executions(
            stateMachineArn=FB_DELIVERY_SFN_ARN,
            statusFilter='RUNNING',
            maxResults=throttle
        )
        num_running = len(sfn_list_response['executions'])
        if num_running < throttle:
            break
        else:
            print(f'Throttling {FB_DELIVERY_SFN_ARN} => {num_running}')
            time.sleep(15)

    def start():
        timed_exc_name = '-'.join([exc_name, str(int(time.time()))])
        print(f'Starting {FB_DELIVERY_SFN_ARN} execution {timed_exc_name}')
        sfn_client.start_execution(
            stateMachineArn=FB_DELIVERY_SFN_ARN,
            name=timed_exc_name,
            input=json.dumps(payload)
        )

    try:
        start()
    except ClientError as e:
        # name collision delivering same sound recording at same exact time
        if e.response['Error']['Code'] == 'ExecutionAlreadyExists':
            time.sleep(1)
            start()
        else:
            raise e


def list_files(sub_dir):
    """List file objs in subdir."""
    ls_dir = f'{PREFIX}{sub_dir}/'
    print(f'Listing files at {ls_dir}')

    list_response = s3_client.list_objects(
        Bucket=BUCKET,
        Prefix=ls_dir
    )
    if list_response['ResponseMetadata']['HTTPStatusCode'] == 200:
        if 'Contents' not in list_response:
            return []
        return [
            x
            for x in list_response['Contents']
            if x['Size'] > 0
        ]
    else:
        print(list_response)
        raise Exception('Unable to list files!')


def load_file(filename, subdir):
    """Load bytes into memory from file in subdir."""
    target_key = f'{PREFIX}{subdir}/{filename}'
    print(f'Loading file at {target_key}')

    get_response = s3_client.get_object(
        Bucket=BUCKET,
        Key=target_key
    )
    if get_response['ResponseMetadata']['HTTPStatusCode'] == 200:
        return get_response['Body'].read().decode('utf-8-sig')
    else:
        print(get_response)
        raise Exception('Unable to load file!')


def move_file(filename, from_subdir, to_subdir):
    """Move file between two sub dirs."""
    old_loc = f'{PREFIX}{from_subdir}/{filename}'
    new_loc = f'{PREFIX}{to_subdir}/{filename}'
    print(f'Moving {old_loc} to {new_loc}')

    copy_response = s3_client.copy_object(
        CopySource={
            'Bucket': BUCKET,
            'Key': old_loc
        },
        Bucket=BUCKET,
        Key=new_loc
    )
    if copy_response['ResponseMetadata']['HTTPStatusCode'] == 200:
        delete_response = s3_client.delete_object(
            Bucket=BUCKET,
            Key=old_loc
        )
        if delete_response['ResponseMetadata']['HTTPStatusCode'] != 204:
            print(delete_response)
            raise Exception('Unable to delete file!')
    else:
        print(copy_response)
        raise Exception('Unable to copy file!')


def invoke_lambda_compilation(osr_id):
    """Invoke lambda for missing compiled version."""
    print(f'Invoking {ADD_VERSION_LAMBDA_NAME} for OSR ID: {osr_id}')
    osr_payload = {
        'label': 'OrchardSoundRecording',
        'id': osr_id
    }

    try:
        lambda_response = lambda_client.invoke(
            FunctionName=ADD_VERSION_LAMBDA_NAME,
            Payload=json.dumps(osr_payload).encode('utf-8')
        )
    except ConnectionClosedError as connection_e:
        raise RetryableException(f'Error Connection close: {connection_e}')

    payload = json.loads(lambda_response['Payload'].read().decode('utf-8'))
    if 'errorMessage' in payload:
        print(payload)
        raise RetryableException(f'Add version invoke failed for OSR ID: {osr_id}')  # noqa:E501

    print('Successful compilation...')
    return payload['results']


if __name__ == '__main__':
    main()