"""Script to trigger redelivery of OSRs based on input CSV file."""

import csv
import json
import sys
import time

import boto3
from botocore.exceptions import ClientError

ACCOUNT_ID = '437795906767'


def extract_latest_version(osr_id, env):
    s3_client = boto3.client('s3', region_name='us-east-1')
    response = s3_client.head_object(
        Bucket=f'{env}-sr-versions-bucket',
        Key=f'{osr_id}',
        ExpectedBucketOwner=ACCOUNT_ID
    )
    version_id = response['VersionId']
    return version_id


def main():
    """Entrypoint."""
    filename = sys.argv[1]
    env = sys.argv[2]
    delivery_type = sys.argv[3]
    after = sys.argv[4] if len(sys.argv) > 4 else None
    max_running = 15

    if delivery_type not in ['FULL_DELIVERY', 'METADATA_UPDATE', 'TAKEDOWN_DELIVERY']:  # noqa:E501
        print("Invalid execution_type! Must be 'FULL_DELIVERY', 'METADATA_UPDATE' or 'TAKEDOWN_DELIVERY'.")  # noqa:E501
        exit(1)

    sfn_name = f'{env}-sr-delivery-tiktok-sfn'  # noqa:E501
    sfn_arn = f'arn:aws:states:us-east-1:437795906767:stateMachine:{sfn_name}'  # noqa:E501
    sfn_client = boto3.client('stepfunctions', region_name='us-east-1')

    # read in CSV file and validate
    osrs = []
    with open(filename, 'r', encoding='utf-8-sig') as f:
        reader = csv.DictReader(f)
        if set(reader.fieldnames) != set(['SOUND_RECORDING_ID']):  # noqa:E501
            print('Invalid input CSV column format!')
            exit(1)
        osrs = [x for x in reader]

    process = False

    num_processed = 0
    total_osrs = len(osrs)
    last_success_osr_id = ""
    for osr in osrs:
        osr_id = osr['SOUND_RECORDING_ID']

        # skip if processing only after specific OSR
        if not process and after:
            num_processed += 1
            if osr_id == after:
                process = True
            continue

        version_id = None
        try:
            version_id = extract_latest_version(osr_id, env)
        except Exception as e:
            if e.response['Error']['Code'] == '404':
                print(f'no such key for osr {osr_id}')
            elif e.response['Error']['Code'] == '400':  # Bad Request, maybe token expired -> renew token and try again  # noqa:E501
                return {'status': 'error', 'message': 'Bad Request', 'last_success_osr_id': last_success_osr_id}  # noqa:E501
            else:
                print(f'Error extracting latest version for OSR {osr_id}: {e}')
            continue
        if not version_id:
            print(f'No version found for OSR {osr_id}, skipping...')
            continue

        is_full_delivery = delivery_type == 'FULL_DELIVERY'

        # throttle based on num running state machines
        while True:
            try:
                sfn_list_response = sfn_client.list_executions(
                    stateMachineArn=sfn_arn,
                    statusFilter='RUNNING',
                    maxResults=max_running
                )
                num_running = len(sfn_list_response['executions'])
                if num_running >= max_running:
                    print("SLEEP")
                    time.sleep(5)
                else:
                    break
            except ClientError as e:
                if e.response['Error']['Code'] == 'ThrottlingException':
                    print(f'ThrottlingException for osr {version_id}')
                    time.sleep(60)
                    break
                elif e.response['Error']['Code'] == 'ExpiredTokenException':
                    return {'status': 'error', 'message': 'AWS credentials have expired', 'last_success_osr_id': last_success_osr_id}  # noqa:E501
                else:
                    raise e

        try:

            # kick-off new execution
            payload = {
                'sound_recording': {
                    'id': osr_id,
                    'version': version_id
                },
                'upload_asset': is_full_delivery,
                'delivery_type': delivery_type
            }
            sfn_client.start_execution(
                stateMachineArn=sfn_arn,
                name='-'.join([version_id, 'manual-delivery', str(int(time.time()))]),  # noqa:E501
                input=json.dumps(payload)
            )
            num_processed += 1
            percent_processed = (num_processed / total_osrs) * 100
            last_success_osr_id = osr_id
            print(f'Processed {osr_id} {num_processed} / {total_osrs} : {percent_processed}%')  # noqa:E501
        except ClientError as e:
            if e.response['Error']['Code'] == 'ExpiredTokenException':
                return {'status': 'error', 'message': 'AWS credentials have expired', 'last_success_osr_id': last_success_osr_id}  # noqa:E501
            print(f'Error processing OSR {osr_id} {version_id}: {e}')
            continue
    return {'status': 'success', 'message': f'Processed {num_processed} OSRs of {total_osrs}', 'last_success_osr_id': last_success_osr_id}  # noqa:E501


if __name__ == '__main__':
    result = main()
    print(result)  # Output result for Jenkins to parse