#!python

import sys,getopt
import os
import json
import shlex, subprocess


def main():

    if not config_file_exit():
        create_config_file()

    refresh_credentials()

    try:
        opts,args = getopt.getopt(sys.argv[1:],"hiuts",["help","train", "upload","infer","setup"])
    except:
        print(err)
        usage()
        sys.exit(2)

    for option,arg in opts:
        if option in ["-t","--train"]:
            cmd = create_train_cmd()
            response = exec_cmd(cmd)
            print(response)
        
        elif option in ["-u", "--upload"]:
            upload_cmd = '''curl -X POST https://w0xf0pfgh8.execute-api.us-east-1.amazonaws.com/dev/upload'''
            response = exec_cmd(upload_cmd)
            timestamp = json.loads(response).get('epoch')       

            config_file = open(".config","a")
            config_file.write(timestamp)
            config_file.close()
            print("Timestamp of your checkpoint is {0}".format(timestamp))

        elif option in ["-i", "--infer"]:
            cmd = create_infer_cmd()
            response = exec_cmd(cmd)
            print(response)

        elif option in ["-s", "--setup"]
            print(arg)
        elif option in ["-h","--help"]:
            usage()
            sys.exit()
        else:
            sys.exit()

def exec_cmd(cmd):
    params = shlex.split(cmd)
    process = subprocess.Popen(params, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    response, stderr = process.communicate()
    return response

def create_train_cmd():
    timestamp = get_timestamp()
    train_cmd = """curl -X POST https://w0xf0pfgh8.execute-api.us-east-1.amazonaws.com/dev/train -d""" + """ '{"epoch":\"""" + timestamp + """"}' """
    return train_cmd

def create_infer_cmd():
    timestamp = get_timestamp()
    infer_cmd = ''' 
    curl -X POST https://w0xf0pfgh8.execute-api.us-east-1.amazonaws.com/dev/infer -d '{"epoch":"''' + timestamp + '''", "input": {"age": ["34"], "workclass": ["Private"], "fnlwgt": ["357145"], "education": ["Bachelors"], "education_num": ["13"], "marital_status": ["Married-civ-spouse"], "occupation": ["Prof-specialty"], "relationship": ["Wife"], "race": ["White"], "gender": ["Female"], "capital_gain": ["0"], "capital_loss": ["0"], "hours_per_week": ["50"], "native_country": ["United-States"], "income_bracket": [">50K"]}}'
     '''
    return infer_cmd

def refresh_credentials():
    refresh_cred_cmd = "aws-okta write-to-credentials geekon/dev /Users/" + os.getenv('USER') + "/.aws/credentials"
    process = subprocess.Popen(refresh_cred_cmd.split(), stdout=subprocess.PIPE)
    output, error = process.communicate()

def config_file_exit():
    return os.path.exists(".config")

def create_config_file():
    f = open(".config","w+")
    f.close()

def get_timestamp():
    f = open(".config","r")
    timestamp = f.read()
    return timestamp

def usage():
    print(''' This script is a utility tool for geekon project which employs AWS lambdas for computation purposes.
    If you want to go through full process firstly you should upload data into S# using -u option, nextly train it with -t and lastly
    receive results with -i option. 
    ''')

if __name__ == "__main__":
    main()