#! /usr/bin/env python
"""
arast-client -- commandline client for Assembly RAST

"""

import os, sys, json, shutil
import argparse
import datetime
import getpass
import logging
import requests
import uuid
import subprocess
import time
from ConfigParser import SafeConfigParser
from pkg_resources import resource_filename

#arast libs
import assembly.asmtypes as asmtypes
import assembly.client as client
import assembly.config as conf
import assembly.auth_token as auth
import traceback


my_version = '0.4.0.0'

# setup option/arg parser
parser = argparse.ArgumentParser(prog='arast', epilog='Use "arast command -h" for more information about a command.')
parser.add_argument('-s', dest='ARAST_URL', help='arast server url')
parser.add_argument('-c', '--config', action="store", help='Specify config file')
parser.add_argument("-v", "--verbose", action="store_true", help="Verbose")
parser.add_argument('--version', action='version', version='AssemblyRAST Client ' + my_version)

subparsers = parser.add_subparsers(dest='command', title='The commands are')

# run -h
p_run = subparsers.add_parser('run', description='Run an Assembly RAST job', help='run job')
data_group = p_run.add_mutually_exclusive_group()
p_run.add_argument("-f", action="append", dest="single", nargs='*', help="specify sequence file(s)")

cmd_group = p_run.add_mutually_exclusive_group()
cmd_group.add_argument("-a", "--assemblers", action="store", dest="assemblers", nargs='*', help="specify assemblers to use. None will invoke automatic mode")
cmd_group.add_argument("-p", "--pipeline", action="append", dest="pipeline", nargs='*', help="invoke a pipeline. None will invoke automatic mode")
cmd_group.add_argument("-r", "--recipe", action="store", dest="recipe", nargs='*', help="invoke a recipe")
cmd_group.add_argument("-w", "--wasp", action="store", dest="wasp", nargs='*', help="invoke a wasp expression")
p_run.add_argument("-m", "--message", action="store", dest="message", help="Attach a description to job")
p_run.add_argument("-q", "--queue", action="store", dest="queue", help=argparse.SUPPRESS)
data_group.add_argument("--data", action="store", dest="data_id", help="Reuse uploaded data")
p_run.add_argument("--pair", action="append", dest="pair", nargs='*', help="Specify a paired-end library and parameters")
p_run.add_argument("--pair_url", action="append", dest="pair_url", nargs='*', help="Specify URLs for a paired-end library and parameters")
p_run.add_argument("--single", action="append", dest="single", nargs='*', help="Specify a single end file and parameters")
p_run.add_argument("--single_url", action="append", dest="single_url", nargs='*', help="Specify a URL for a single end file and parameters")
p_run.add_argument("--reference", action="append", dest="reference", nargs='*', help="specify sequence file(s)")
p_run.add_argument("--reference_url", action="append", dest="reference_url", nargs='*', help="Specify a URL for a reference contig file and parameters")
p_run.add_argument("--curl", action="store_true", help="Use curl for http requests")

# stat -h
p_stat = subparsers.add_parser('stat', description='Query status of running jobs', help='list jobs status')
p_stat.add_argument("-j", "--job", action="store", help="get status of specific job")
p_stat.add_argument("-w", "--watch", action="store_true", help="monitor in realtime")
p_stat.add_argument("-n", dest="stat_n", action="store", default=10, type=int, help="specify number of records to show")
p_stat.add_argument("-l", "--list-data", action="store_true", dest="list_data", help="list data objects")
p_stat.add_argument("--data-json", action="store", dest="data_id", help="print json string for data object")

# avail
p_avail = subparsers.add_parser('avail', description='List available AssemblyRAST modules', help='list available modules')
p_avail.add_argument("-v", "--verbose", action="store_true", help="show module details")

# upload
p_upload = subparsers.add_parser('upload', description='Upload a read set', help='Upload a read library or set of libraries, returns a data ID for future use')
p_upload.add_argument("-f", action="append", dest="single", nargs='*', help="specify sequence file(s)")
p_upload.add_argument("--pair", action="append", dest="pair", nargs='*', help="Specify a paired-end library and parameters")
p_upload.add_argument("--pair_url", action="append", dest="pair_url", nargs='*', help="Specify URLs for a paired-end library and parameters")
p_upload.add_argument("--single", action="append", dest="single", nargs='*', help="Specify a single end file and parameters")
p_upload.add_argument("--single_url", action="append", dest="single_url", nargs='*', help="Specify a URL for a single end file and parameters")
p_upload.add_argument("--reference", action="append", dest="reference", nargs='*', help="specify a reference contig file")
p_upload.add_argument("--reference_url", action="append", dest="reference_url", nargs='*', help="Specify a URL for a reference contig file and parameters")
p_upload.add_argument("-m", "--message", action="store", dest="message", help="Attach a description to job")
p_upload.add_argument("--json", action="store_true", help="Print data info json object to STDOUT")

# kill
p_kill = subparsers.add_parser('kill', description='Send a kill signal to jobs', help='kill jobs')
p_kill.add_argument("-j", "--job", action="store", help="kill specific job")
p_kill.add_argument("-a", "--all", action="store_true", help="kill all user jobs")

# get
p_get = subparsers.add_parser('get', description='Download result data', help='download data')
p_get.add_argument("-j", "--job", action="store", dest="job_id", nargs=1, required=True, help="specify which job data to get")
p_get.add_argument("-a", "--assembly", action="store", nargs='?', default=False, const=True, help="Get assemblies only")
p_get.add_argument("-r", "--report", action="store_true", help="Print assembly report to stdout")
p_get.add_argument("--stdout", action="store_true", help="Print assembly to stdout")
p_get.add_argument("-o", "--outdir", action="store", help="Download to specified dir")
p_get.add_argument("-w", "--wait", action="store_true", help="Wait until job is done")

p_logout = subparsers.add_parser('logout', description='Log out', help='log out')
p_login = subparsers.add_parser('login', description='Force log in', help='log in')
p_login.add_argument("--rast", action="store_true", help="Print assembly to stdout")

def main():
    global aclient
    
    clientlog = logging.getLogger('client')
    clientlog.setLevel(logging.INFO)
    sh = logging.StreamHandler()
    sh.setLevel(logging.DEBUG)
    frmt = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    sh.setFormatter(frmt)
    clientlog.addHandler(sh)

    args = parser.parse_args()
    opt = parser.parse_args()
    options = vars(args)
    
    options['version'] = my_version

    if args.verbose:
        clientlog.setLevel(logging.DEBUG)
        clientlog.debug("Logger Debugging mode")

    #### Get configuration #####
    ARAST_URL = conf.URL
    user_dir = user_data_dir(conf.APPNAME, conf.APPAUTHOR)

    oauth_file = os.path.join(user_dir, conf.OAUTH_FILENAME)
    expiration = conf.OAUTH_EXP_DAYS

    oauth_parser = SafeConfigParser()
    oauth_parser.read(oauth_file)
    reauthorize = True

    if "KB_RUNNING_IN_IRIS" in os.environ:
        if args.command == 'logout' or args.command == 'login':
            print "Please use the IRIS controls to log in/out"
            sys.exit()
        if "KB_AUTH_TOKEN" in os.environ and "KB_AUTH_USER_ID" in os.environ and \
                len(os.environ["KB_AUTH_USER_ID"]) > 0 and \
                len(os.environ["KB_AUTH_TOKEN"]) > 0 :
            a_user = os.environ["KB_AUTH_USER_ID"]
            a_token = os.environ["KB_AUTH_TOKEN"]
        else:
            print "Please authenticate with KBase credentials"
            sys.exit()

    else:
        if args.command == 'logout' or args.command == 'login':
            try:
                os.remove(oauth_file)
            except:
                pass
            if args.command == 'logout':
                print >> sys.stderr, '[x] Logged out'
                sys.exit()
    
        # Check if user file exists
        if os.path.exists(oauth_file):
            token_date_str = oauth_parser.get('auth', 'token_date')
            tdate = datetime.datetime.strptime(token_date_str, '%Y-%m-%d').date()
            cdate = datetime.date.today()

            if (cdate - tdate).days > expiration:
                reauthorize = True
            else:
                reauthorize = False

        if not reauthorize:
            a_user = oauth_parser.get('auth', 'user')
            a_token = oauth_parser.get('auth', 'token')
            # print >> sys.stderr, "Logged in as: {}".format(a_user)
        else:
            if args.rast:
                print("Please authenticate with RAST credentials")
                a_user = raw_input("RAST Login: ")
                a_pass = getpass.getpass(prompt="RAST Password: ")
                globus_map = auth.get_token(a_user, a_pass, auth_svc=auth.RAST_URL)
                a_token = globus_map['access_token']
                a_user = '{}_rast'.format(a_user)
            else:
                print("Please authenticate with KBase credentials")
                a_user = raw_input("KBase Login: ")
                a_pass = getpass.getpass(prompt="KBase Password: ")
                globus_map = auth.get_token(a_user, a_pass)
                a_token = globus_map['access_token']
                try:
                    os.makedirs(user_dir)
                except:
                    pass
            uparse = SafeConfigParser()
            uparse.add_section('auth')
            uparse.set('auth', 'user', a_user)
            uparse.set('auth', 'token', a_token)
            uparse.set('auth', 'token_date', str(datetime.date.today()))
            uparse.write(open(oauth_file, 'wb'))

    if args.command == 'login':
        sys.exit()
    
    if args.ARAST_URL:
        ARAST_URL = args.ARAST_URL

    aclient = client.Client(ARAST_URL, a_user, a_token)

        
    res_ids = []
    file_sizes = []
    file_list = []


    curl = False
    try:
        curl = args.curl
    except:
        pass


    # Format into separate pipelines
    if args.command == "run" or args.command == "upload":
        if args.command == "run":
            if args.assemblers:
                args.pipeline = [(" ".join(args.assemblers))]

            if not args.pipeline: # auto
                args.pipeline = 'auto'

            if not args.pipeline:
                parser.print_usage()
                sys.exit()

            if not (args.data_id or 
                    args.pair or args.pair_url or
                    args.single or args.single_url):
                parser.print_usage()
                sys.exit()

        if args.command == "upload" and not (args.pair or args.single):
            parser.print_usage()
            sys.exit()

        files = []
        adata = client.AssemblyData()
        
        ##### Parse args and create AssemblyData dict #####
        try:
            has_data_id = args.data_id
        except:
            has_data_id = False
        if not has_data_id:
            all_lists = [args.pair, args.pair_url, args.single, args.single_url, args.reference, args.reference_url]
            file_lists = []
            for l in all_lists:
                if l is None:
                    file_lists.append([])
                else:
                    file_lists.append(l)
                   
            all_types = ['paired', 'paired_url', 'single', 'single_url', 'reference', 'reference_url']
            for f_list, f_type in zip(file_lists, all_types):
                for ls in f_list:
                    f_infos = []
                    f_set_args = {}
                    for word in ls:
                        if not (os.path.isfile(word) or '=' in word or is_valid_url(word)):
                            raise Exception('{} is not valid input!'.format(word))
                    for word in ls:
                        if os.path.isfile(word):
                            f_info = aclient.upload_data_file_info(word, curl=curl)
                            f_infos.append(f_info)
                        elif '=' in word:
                            kv = word.split('=')
                            f_set_args[kv[0]] = kv[1]
                        elif is_valid_url(word):
                            f_info = asmtypes.FileInfo(direct_url=word)
                            f_infos.append(f_info)
                    f_set = asmtypes.FileSet(f_type, f_infos, **f_set_args)
                    adata.add_set(f_set)

        arast_msg = dict((k, options[k]) for k in ['pipeline', 'data_id', 'message', 'queue', 'version', 'recipe', 'wasp'] if k in options)

        arast_msg['assembly_data'] = adata
        arast_msg['client'] = 'CLI'

        ##### Send message to Arast Server #####
        payload = json.dumps(arast_msg, sort_keys=True)
        clientlog.debug(" [x] Sending message: %r" % (payload))

        if args.command == "run":
            response = aclient.submit_job(payload)
            print 'Job ID: {}'.format(response)
        if args.command == "upload":
            response = aclient.submit_data(payload)
            arast_msg.update(json.loads(response))
            if args.json:
                # print arast_msg
                print payload
            print 'Data ID: {}'.format(arast_msg['data_id'])

    elif args.command == 'stat':
        if args.list_data:
            table = aclient.get_data_list_table(args.stat_n)
            print table
            sys.exit()

        if args.data_id:
            data_json = aclient.get_data_json(args.data_id)
            print data_json
            sys.exit()
        
        # default: print job information
        while True:
            try:
                response = aclient.get_job_status(args.stat_n, args.job)
                if args.watch:
                        os.system('clear')
                print response
                if not args.watch:
                        break
                else:
                    print 'Press CTRL-C to quit.'
                ### Spinner loop
                spinners = ['-', '\\', '|', '/'] 
                sleep_seconds = 25
                spins_per_sec = 4
                for i in range(sleep_seconds * spins_per_sec):
                    os.system('clear')
                    print('[{}] Assembly Service Status').format(spinners[i%4])
                    print response
                    print 'Press CTRL-C to quit.'
                    time.sleep(1.0/spins_per_sec)			
            except KeyboardInterrupt:
                break
                

    elif args.command == 'get':
        aclient.validate_job(args.job_id[0])

        if args.wait:
            stat = aclient.wait_for_job(args.job_id[0])
            if 'FAIL' in stat:
                print stat
                sys.exit()

        if args.report:
            # try:
            report = aclient.get_job_report(args.job_id[0])
            print report
            sys.exit()

        if args.assembly:
            try:
                if type(args.assembly) is str:
                    aclient.get_assemblies(job_id=args.job_id[0], asm_id=args.assembly, stdout=args.stdout, outdir=args.outdir)
                else:
                    aclient.get_assemblies(job_id=args.job_id[0], stdout=args.stdout, outdir=args.outdir)
            except:
                print traceback.format_tb(sys.exc_info()[2])
                print sys.exc_info()
        else:
            try:
                aclient.get_job_data(job_id=args.job_id[0], outdir=args.outdir)
            except:
                print 'Error retrieving job data: {}'.format(args.job_id[0])

    elif args.command == 'avail':
        try:
            mods = json.loads(aclient.get_available_modules())
            mods = sorted(mods, key=lambda mod: mod['module'])
            # print json.dumps(mods, indent=4)
            
            if args.verbose:
                for mod in mods:
                    keys = ('description', 'version', 'stages', 'modules', 'limitations', 'references')
                    if mod['version'] >= '1.0':
                        print '[Module] ' + mod['module']
                        for key in keys:
                            if key in mod.keys():
                                print '  '+key.title()+': '+mod[key]

                        if 'parameters' in mod.keys() :
                            parms = mod['parameters']
                            if len(parms) > 0:
                                print '  Customizable parameters: default (available values)'
                                for parm in sorted(parms, key=lambda p: p[0]):
                                    print '%25s  =  %s' % (parm[0], parm[1])
                        print
            else:
                print '{0:16} {1:35} {2:10}'.format('Module', 'Stages', 'Description')
                print '----------------------------------------------------------------'
                for mod in mods:
                    if mod['version'] >= '1.0':
                        print '{module:16} {stages:35} {description}'.format(**mod)

        except:
            print 'Error getting available modules'

    elif args.command == 'kill':
        print aclient.kill_jobs(args.job)

def is_valid_url(url):
    import re
    regex = re.compile(
        r'^https?://'  # http:// or https://
        r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+[A-Z]{2,6}\.?|'  # domain...
        r'localhost|'  # localhost...
        r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})' # ...or ip
        r'(?::\d+)?'  # optional port
        r'(?:/?|[/?]\S+)$', re.IGNORECASE)
    return url is not None and regex.search(url)

def user_data_dir(appname, appauthor):
     return os.path.expanduser('/'.join(['~', '.config', appname]))

if __name__ == '__main__':
    main()
