import os

# Include the Lambda layer extracted location

import boto3
import time
import datetime
import json
import sys

region = 'eu-west-1'
if len(sys.argv) < 3:
    print('''Usage program.py data_set_id bucket_name
    eg program.py sjhjsgdhsjgdshjdgsh garys_bucket''')
data_set_id = sys.argv[1]
bucket_name = sys.argv[2]
object_key = ''

if os.environ.get('MALWARE_PATH'):
    if isinstance(os.environ['MALWARE_PATH'], str):
        path = os.environ['MALWARE_PATH']   
    else: 
        print(f"Path doesn't appear to be a string")
else: 
    print('You must set the environement variable MALWARE_PATH')
    exit()




def recurse_files(path):
    good_count = 0
    count = 0
    s3 = boto3.resource('s3')
    for dirpath, dirs, files in os.walk(path):
        for file in files:
            s3.Bucket(bucket_name).upload_file(os.path.join(dirpath, file), file)
            message, good = create_job(data_set_id, bucket_name, file, region)
            count = count +1 
            good_count = good_count + good
            print(f"File {os.path.join(dirpath, file)} returned {message}")
    return count, good_count

def create_job(data_set_id, bucket_name, object_key, region):
    # Setup the boto3 clients needed
    dataexchange = boto3.client(
        service_name='dataexchange',
        region_name=region
        )


    # CREATE REVISION under the dataset provided as an environment variable
    current_time_for_creating_revision = datetime.datetime.utcnow().strftime("%d %B %Y %I:%M%p UTC")
    create_revision_response = dataexchange.create_revision(
                                                DataSetId=data_set_id,
                                                Comment='Revision created programmatically on '
                                                + current_time_for_creating_revision)
    revision_id = create_revision_response['Id']

    # CREATE JOB under the revision to import file from S3 to DataExchange
    create_job_s3_import = dataexchange.create_job(
        Type='IMPORT_ASSETS_FROM_S3',
        Details={
            'ImportAssetsFromS3': {
                'DataSetId': data_set_id,
                'RevisionId': revision_id,
                'AssetSources': [
                    {
                        'Bucket': bucket_name,
                        'Key': object_key
                    }

                ]
            }
        }
    )

    # Filter the ID of the Job from the response
    job_id = create_job_s3_import['Id']

    # invoke START JOB on the created job to change it from Waiting to
    # Completed state
    start_created_job = dataexchange.start_job(JobId=job_id)

    # GET JOB details to track the state of the job and wait until it reaches
    # COMPLETED state
    job_status = ''

    while job_status != 'COMPLETED':
        get_job_status = dataexchange.get_job(JobId=job_id)
        job_status = get_job_status['State']
        message = 'Job Status ' + job_status
        print(message)
        good = 1
        if job_status == 'ERROR':
            job_errors = get_job_status['Errors']
            message = 'JobId: {} failed with error:{}'.format(job_id, job_errors)  
            return message, good     
        time.sleep(3)        
    # Finalize revision by invoking UPDATE REVISION
    current_time_for_finalize_revision = datetime.datetime.utcnow().strftime("%d %B %Y %I:%M%p UTC")
    #print(current_time_for_finalize_revision)
    finalize_revision = dataexchange.update_revision(DataSetId=data_set_id, 
                                                    RevisionId=revision_id, 
                                                    Finalized=True,
                                                    Comment='Revision finalized programmatically on ' + current_time_for_finalize_revision)

    return message, good

if __name__ == '__main__':
    total, good = recurse_files(path)
    print(f'Completed {total} with {good} judged as OK')
