#!/usr/bin/env python

# https://github.com/ross/requests-futures
# http://docs.python-requests.org/en/master/user/quickstart/#more-complicated-post-requests

import sys, os, uuid, optparse, requests, json, time
#from requests_futures.sessions import FuturesSession

#### NN14 ####
SERVICE_URL = "http://nn14.galaxyproject.org:8080/";
#service_url = "http://127.0.0.1:8082/";
QUERY_URL = SERVICE_URL+"tree/<tree_id>/query";
STATUS_URL = SERVICE_URL+"status/<query_id>";
##############
# query delay in seconds
QUERY_DELAY = 30;
##############

__version__ = "1.0.0";
VALID_CHARS = '.-()[]0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ '
# in the case of collections, exitcodes equal to 0 and 1 are not considered errors
ERR_EXIT_CODE = 2;
OK_EXIT_CODE = 0;

def raiseException( exitcode, message, output_dir_path, errorfilename ):
    errorfilepath = os.path.join(output_dir_path, errorfilename+"_txt");
    with open(errorfilepath, 'w') as out:
        out.write(message);
    sys.exit(exitcode);

def query_request( options, args, payload ):
    output_dir_path = options.outputdir;
    # add additional parameters to the payload
    #payload["tree_id"] = str(options.treeid);
    payload["search_mode"] = str(options.search);
    payload["exact_algorithm"] = int(options.exact);
    payload["search_threshold"] = float(options.sthreshold);
    payload["sort"] = int(options.sortcontrol);
    # set the content type to application/json
    headers = {'Content-type': 'application/json'};

    # create a session
    session = requests.Session();
    # make a synchronous post request to the query route
    req = session.post(QUERY_URL.replace("<tree_id>", str(options.treeid)), headers=headers, json=payload);
    resp_code = req.status_code;
    #print(str(req.content)+"\n\n");
    if resp_code == requests.codes.ok:
        resp_content = str(req.content);
        # convert out to json
        json_content = json.loads(resp_content);
        # retrieve query id
        query_id = json_content['query_id'];
        query_processed = False;
        # results json content
        json_status_content = None;
        while query_processed is False:
            # create a new session
            session = requests.Session();
            # make a synchronous get request to the status route
            status_query_url = STATUS_URL.replace("<query_id>", query_id);
            status_req = session.get(status_query_url);
            status_resp_content = str(status_req.content);
            #print(status_resp_content+"\n\n");
            # convert out to json
            json_status_content = json.loads(status_resp_content);
            # take a look at the state
            # state attribute is always available
            if json_status_content['state'] == 'SUCCESS':
                query_processed = True;
                break;
            elif json_status_content['state'] in ['FAILURE', 'REVOKED']:
                return raiseException( ERR_EXIT_CODE, "Query ID: "+str(query_id)+"\nQuery status: "+str(json_status_content['state']), output_dir_path, str(options.errorfile) );
            else:
                time.sleep(QUERY_DELAY); # in seconds
        
        out_file_format = "tabular";
        for block in json_status_content['results']:
            seq_id = block['sequence_id'];
            # put response block in the output collection
            output_file_path = os.path.join(output_dir_path, seq_id + "_" + out_file_format);
            accessions_list = "";
            hits_block = block['hits'];
            accessions_dict = { };
            is_sabutan = False;
            for hit in hits_block:
                if type(hit) is dict: # sabutan
                    #accessions_list = accessions_list + str(hit['accession_number']) + "\t" + str(hit['score']) + "\n";
                    accession_number = hit['accession_number'];
                    #------------
                    #score = hit['score'];
                    #score_split = score.split("/");
                    #accessions_dict[accession_number] = "{0:.6f}".format(float(score_split[0])/float(score_split[1]));
                    #------------
                    fraction = hit['fraction'];
                    score = hit['score'];
                    accession_scores = {
                        "fraction": str(fraction),
                        "score": float(score)
                    }
                    accessions_dict[accession_number] = accession_scores;
                    is_sabutan = True;
                else: # all-some
                    accessions_list = accessions_list + str(hit) + "\n";
            if is_sabutan:
                sorted_accessions = sorted(accessions_dict, key=lambda i: float(accessions_dict[i]["score"]), reverse=True);
                for acc in sorted_accessions:
                    accessions_list = accessions_list + str(acc) + "\t" + str(accessions_dict[acc]["fraction"]) + "\t" + str(accessions_dict[acc]["score"]) + "\n";
            with open(output_file_path, 'w') as out:
                out.write(accessions_list.strip());
        return sys.exit(OK_EXIT_CODE);
    else:
        return raiseException( ERR_EXIT_CODE, "Unable to query the remote server. Please try again in a while.", output_dir_path, str(options.errorfile) );

def query( options, args ):
    output_dir_path = options.outputdir;
    multiple_data = {};
    comma_sep_file_paths = options.files;
    #print("files: "+str(comma_sep_file_paths)+" - "+str(type(comma_sep_file_paths)));
    # check if options.files contains at least one file path
    if comma_sep_file_paths is not None:
        # split file paths
        file_paths = comma_sep_file_paths.split(",");
        # split file names
        comma_sep_file_names = str(options.names);
        #print("names: "+str(comma_sep_file_names));
        file_names = comma_sep_file_names.split(",");
        for idx, file_path in enumerate(file_paths):
            #file_name = file_names[idx];
            with open(file_path, 'r') as content_file:
                for line in content_file:
                    if line.strip() != "":
                        line_split = line.strip().split("\t"); # split on tab
                        if len(line_split) == 2: # 0:id , 1:seq , otherwise skip line
                            seq_id = line_split[0];
                            # fix seq_id using valid chars only
                            seq_id = ''.join(e for e in seq_id if e in VALID_CHARS)
                            seq_text = line_split[1];
                            if seq_id in multiple_data:
                                return raiseException( ERR_EXIT_CODE, "Error: the id '"+seq_id+"' is duplicated", output_dir_path, str(options.errorfile) );
                            multiple_data[seq_id] = seq_text;
        if len(multiple_data) > 0:
            return query_request( options, args,  multiple_data );
            #return echo( options, args );
        else:
            return raiseException( ERR_EXIT_CODE, "An error has occurred. Please be sure that your input files are valid.", output_dir_path, str(options.errorfile) );
    else:
        # try with the sequence in --sequence
        text_content = options.sequences;
        #print("sequences: "+text_content);
        # check if options.sequences contains a list of sequences (one for each row)
        if text_content is not None:
            text_content = str(text_content);
            if text_content.strip():
                # populate a dictionary with the files containing the sequences to query
                text_content = text_content.strip().split("__cn__"); # split on new line
                for line in text_content:
                    if line.strip() != "":
                        line_split = line.strip().split("__tc__"); # split on tab
                        if len(line_split) == 2: # 0:id , 1:seq , otherwise skip line
                            seq_id = line_split[0];
                            # fix seq_id using valid chars only
                            seq_id = ''.join(e for e in seq_id if e in VALID_CHARS)
                            seq_text = line_split[1];
                            if seq_id in multiple_data:
                                return raiseException( ERR_EXIT_CODE, "Error: the id '"+seq_id+"' is duplicated", output_dir_path, str(options.errorfile) );
                            multiple_data[seq_id] = seq_text;
                if len(multiple_data) > 0:
                    return query_request( options, args, multiple_data );
                    #return echo( options, args );
                else:
                    return raiseException( ERR_EXIT_CODE, "An error has occurred. Please be sure that your input files are valid.", output_dir_path, str(options.errorfile) );
            else:
                return raiseException( ERR_EXIT_CODE, "You have to insert at least one row formatted as a tab delimited (ID, SEQUENCE) couple", output_dir_path, str(options.errorfile) );
    return ERR_EXIT_CODE;

def __main__():
    # Parse the command line options
    usage = "Usage: query.py --files comma_sep_file_paths --names comma_seq_file_names --sequences sequences_text --search search_mode --exact exact_alg --sthreshold threshold --outputdir output_dir_path";
    parser = optparse.OptionParser(usage = usage);
    parser.add_option("-v", "--version", action="store_true", dest="version",
                    default=False, help="display version and exit")
    parser.add_option("-f", "--files", type="string",
                    action="store", dest="files", help="comma separated files path");
    parser.add_option("-n", "--names", type="string",
                    action="store", dest="names", help="comma separated names associated to the files specified in --files");
    parser.add_option("-s", "--sequences", type="string",
                    action="store", dest="sequences", help="contains a list of sequences (one for each row)");
    parser.add_option("-a", "--fasta", type="string",
                    action="store", dest="fasta", help="contains the content of a fasta file");
    parser.add_option("-x", "--search", type="string", default="rrr",
                    action="store", dest="search", help="search mode");
    parser.add_option("-e", "--exact", type="int", default=0,
                    action="store", dest="exact", help="exact algorithm (required if search is 1 only)");
    parser.add_option("-k", "--tree", type="string", default=0,
                    action="store", dest="treeid", help="the id of the tree that will be queried");
    parser.add_option("-t", "--sthreshold", type="float",
                    action="store", dest="sthreshold", help="threshold applied to the search algrithm");
    parser.add_option("-z", "--sort", type="int", default=1,
                    action="store", dest="sortcontrol", help="boolean required to sort the result");
    parser.add_option("-o", "--outputdir", type="string", default="output",
                    action="store", dest="outputdir", help="output directory (collection) path");
    parser.add_option("-r", "--errorfile", type="string", default="error_txt",
                    action="store", dest="errorfile", help="error file name containing error messages");

    # TEST
    #sequences = 'NM_001169378.2__tc__atttcggatgctttggagggaggaactctagtgctgcattgattggggcgtgtgttaatgatattcccagttcgcatggcgagcatcgattcctggtacgtatgtgggccccttgactcccacttatcgcacttgtcgttcgcaatttgcatgaattccgcttcgtctgaaacgcacttgcgccagacttctccggctggtctgatctggtctgtgatccggtctggtggggcgccagttgcgtttcgagctcatcaccagtcactccgcagtcgcattctgccagaggtctccgatcaagagcgcttctccattcgagattcaaacgcagcgcggtctgacgccgccacatcgagtgaaatccatatcgatggccacattcacacaggacgagatcgacttcctgcgcagccatggcaacgagctgtgtgccaagacctggctgggattgtgggatccgaagcgggctgtgcaccagcaggagcagcgcgaactgatgatggacaagtatgagcggaagcgatactacctggagccggccagtcctcttaagtcgctggccaatgcggtcaacctgaagtcgtctgctccggcgacgaaccacactcagaatggccaccaaaatgggtatgccagcatccatttgacgcctcctgctgcccagcggacctcggccaatggattgcagaaggtggccaactcgtcgagtaactcttctggaaagacctcatcctcgatcagtaggccacactataatcaccagaacaacagccaaaacaacaatcacgatgcctttggcctgggtggcggattgagcagcctgaacagcgccggttccacatccactggagctctttccgacaccagcagttgtgctagcaatggcttcggtgcggactgcgactttgtggctgactttggctcggccaacattttcgacgccacatcggcgcgttccacaggatcgccggcggtgtcgtccgtgtcctcagtgggttccagcaatggctacgccaaggtgcagcccatccgggcagctcatctccagcagcaacagcagttgcagcagcagctgcatcagcagcagctcctcaatggcaatggtcatcagggcactgagaactttgccgacttcgatcacgctcccatctacaatgcagtggctccaccgacttttaacgattggatcagcgactggagcaggcggggcttccacgatcccttcgacgattgcgatgactcgccaccaggtgcccgccctccagcacctgcgccagctcctgctcaagttcccgcagtatcatcaccattgccaaccgtccgagaagaaccagagcttgcgtggaatttttgggaggacgagatgcgaatagaggcgcaggaaaaggagtcccaaactaaacagccggagttgggctactccttttcgattagtactactacgcccctttccccttcgaatcccttcctgccctaccttgtcagtgaggagcagcatcgaaatcatccagagaagccctccttttcgtattcgttgttcagctccatatcaaatagttcgcaagaagatcaggcggatgatcatgagatgaatgttttaaatgccaatttccatgatttctttacgtggagtgctcccttgcagaacggccatacgaccagtccgcccaagggcggaaatgcagcgatggcgcccagtgaggatcgatatgccgctcttaaggatctcgacgagcagctgcgagaactgaaggccagcgaaagcgccacagagacgcccacgcccaccagtggcaatgttcaggccacagatgcctttggtggagccctcaacaacaatccaaatcccttcaagggccagcaacagcagcagctcagcagccatgtggtgaatccattccagcagcagcaacagcagcagcaccagcagaatctctatggccagttgacgctcataccaaatgcctacggcagcagttcccagcagcagatggggcaccatctcctccagcagcagcagcagcaacagcagagcttcttcaacttcaacaacaacgggttcgccatctcgcagggtctgcccaacggctgcggcttcggcagcatgcaacccgctcctgtgatggccaacaatccctttgcagccagcggcgccatgaacaccaacaatccattcttatgagactcaacccgggagaatccgcctcgcgccacctggcagaggcgctgagccagcgaacaaagagcagacgcggaggaaccgaaccgaaattagtccattttactaacaatagcgttaatctatgtatacataatgcacgccggagagcactctttgtgtacatagcccaaatatgtacacccgaaaggctccacgctgacgctagtcctcgcggatggcggaggcggactggggcgttgatatattcttttacatggtaactctactctaacgtttacggatacggatatttgtatttgccgtttgccctagaactctatacttgtactaagcgcccatgaacacttcatccactaacatagctactaatcctcatcctagtggaggatgcagttggtccagacactctgttatttgttttatccatcctcgtacttgtctttgtcccatttagcactttcgttgcggataagaactttgtcagttattgattgtgtggccttaataagattataaaactaaatattataacgtacgactatacatatacggatacagatacagattcagacacagttagtacagatacagatatacatatacgcttttgtacctaatgaattgcttcttgtttccattgctaatcatctgcttttcgtgtgctaattttatacactagtacgtgcgatatcggccgtgcagatagattgctcagctcgcgagtcaagcctcttttggttgcacccacggcagacatttgtacatatactgtctgattgtaagcctcgtgtaatacctccattaacaccactcccccaccacccatccatcgaaccccgaatccatgactcaattcactgctcacatgtccatgcccatgccttaacgtgtcaaacattatcgaagccttaaagttatttaaaactacgaaatttcaataaaaacaaataagaacgctatc';
    #(options, args) = parser.parse_args(['-x', 'rrr', '-t', 0.5, '-s', sequences, '-o', 'collection_content']);

    (options, args) = parser.parse_args();
    if options.version:
        print __version__;
    else:
        # create output dir (collection)
        output_dir_path = options.outputdir;
        if not os.path.exists(output_dir_path):
            os.makedirs(output_dir_path);

        return query( options, args );

if __name__ == "__main__": __main__()