#!/usr/bin/env python

import argparse
import os
import sys
import json
import time
import re
import traceback
from urllib import urlencode
from urllib2 import Request, urlopen, HTTPError

if sys.version_info.major != 2:
    raise Exception('This tool supports only python2')


argparser = argparse.ArgumentParser(description = 'Manage Google Oauth 2.0 authentication for Bravo API and do querying from command line.')
commands = argparser.add_subparsers(dest = 'command', title = 'Commands')
login_command = commands.add_parser('login', help = 'Authorize access to Bravo API.')
token_command = commands.add_parser('print-access-token', help = 'Display access token for Bravo API.')
revoke_command = commands.add_parser('revoke', help = 'Revoke all access tokens for Bravo API.')
query_region_command = commands.add_parser('query-region', help = 'Query chromosomal region.')
query_gene_command = commands.add_parser('query-gene', help = 'Query by gene name or gene identifier.')
query_variant_command = commands.add_parser('query-variant', help = 'Query variant by variant identifier or by chromosome name and chromosomal position.',
        description = 'Query variant by identifier CHROM-POS-REF-ALT, or by chromosome name and chromosomal position.')
query_meta_command = commands.add_parser('query-meta', help = 'Query version and data description.')
annotate_command = commands.add_parser('annotate', help = 'Annotate input VCF.',
        description = 'Uncompressed input VCF must be streamed to standard input. Uncompressed output VCF is streamed to standard output. Multi-allelic variant records in input VCF must be split into multiple bi-allelic variant records.')

query_region_command.add_argument('-c', '--chromosome', metavar = 'name', type = str, required = True, dest = 'chromosome', help = 'Chromosome name.')
query_region_command.add_argument('-s', '--start', metavar = 'base-pair', type = int, required = True, dest = 'start', help = 'Start position.')
query_region_command.add_argument('-e', '--end', metavar = 'base-pair', type = int, required = True, dest = 'end', help = 'End position.')
query_region_command.add_argument('-o', '--output', required = False, choices = ['json', 'vcf'], default = 'json', dest = 'format', help = 'Output format.')
query_region_command.add_argument('-f', '--filter', metavar = 'expression', required = False, type = str, dest = 'filter', help = 'Filtering expression.')

query_gene_command.add_argument('-n', '--name', metavar = 'name', type = str, required = True, dest = 'gene', help = 'Gene name or gene identifier.')
query_gene_command.add_argument('-o', '--output', required = False, choices = ['json', 'vcf'], default = 'json', dest = 'format', help = 'Output format.')
query_gene_command.add_argument('-f', '--filter', metavar = 'expression', required = False, type = str, dest = 'filter', help = 'Filtering expression.')

query_variant_command.add_argument('-v', '--variant', metavar = 'chrom-pos-ref-alt/rs#', type = str, dest = 'variant_id', help = 'Variant identifier CHROM-POS-REF-ALT or rs#.')
query_variant_command.add_argument('-c', '--chromosome', metavar = 'name', type = str, dest = 'chromosome', help = 'Chromosome name.')
query_variant_command.add_argument('-p', '--position', metavar = 'base-pair', type = int, dest = 'position', help = 'Position.')
query_variant_command.add_argument('-o', '--output', required = False, choices = ['json', 'vcf'], default = 'json', dest = 'format', help = 'Output format.')

annotate_command.add_argument('-f', '--filter', metavar = 'expression', required = False, type = str, dest = 'filter', help = 'Filtering expression.') 

USER_HOME = os.path.expanduser("~")
BRAVO_DIR = '.bravo'
BRAVO_CREDSTORE = 'credstore'
BRAVO_API_VERSION = 'v1'
BRAVO_AUTH_API = 'https://bravo.sph.umich.edu/api/{}/auth/auth'.format(BRAVO_API_VERSION)
BRAVO_TOKEN_API = 'https://bravo.sph.umich.edu/api/{}/auth/token'.format(BRAVO_API_VERSION)
BRAVO_REVOKE_API = 'https://bravo.sph.umich.edu/api/{}/auth/revoke'.format(BRAVO_API_VERSION)
BRAVO_IP_API = 'https://bravo.sph.umich.edu/api/{}/auth/ip'.format(BRAVO_API_VERSION)


class BravoException(Exception):
    def __init__(self, message):
        self.message = message
    def __str__(self):
        return self.message


class requests(object): # this implements the parts we need of the real `requests` module
    @staticmethod
    def get(url, headers={}, params=None):
        if params: url += '?' + urlencode(params)
        r = Request(url, headers=headers)
        return _requests_response(urlopen(r))
    @staticmethod
    def post(url, headers={}, data=None):
        if data is not None and isinstance(data, dict): data = urlencode(data)
        r = Request(url, headers=headers, data=data)
        return _requests_response(urlopen(r))
class _requests_response(object):
    def __init__(self, response):
        self.status_code = response.getcode()
        self._json = json.load(response.fp)
    def json(self):
        return self._json
class _requests_exceptions(object): pass
requests.exceptions = _requests_exceptions
requests.exceptions.RequestException = HTTPError


def credstore_exists():
    return os.path.isfile(os.path.join(USER_HOME, BRAVO_DIR, BRAVO_CREDSTORE))


def create_credstore():
    p = os.path.join(USER_HOME, BRAVO_DIR)
    if not os.path.isdir(p):
        os.mkdir(p, 0700)
    p = os.path.join(p, BRAVO_CREDSTORE)
    if not os.path.isfile(p):
        open(p, 'a').close()
        os.chmod(p, 0600)


def read_credstore():
    required = ['access_token', 'created', 'token_type', 'revoked']
    with open(os.path.join(USER_HOME, BRAVO_DIR, BRAVO_CREDSTORE), 'r') as credstore_file:
        credstore = json.load(credstore_file)
        active_credentials = credstore.get('active', None)
        if active_credentials is None or not isinstance(active_credentials, basestring):
            raise BravoException('Invalid or outdated credentials store. You may need to run login.')
        all_credentials = credstore.get('all', None)
        if all_credentials is None or not isinstance(all_credentials, dict):
            raise BravoException('Invalid or outdated credentials store. You may need to run login.')
        if active_credentials not in all_credentials:
            raise BravoException('Invalid or outdated credentials store. You may need to run login.')
        for ip, credentials in all_credentials.iteritems():
            if not isinstance(credentials, dict):
                raise BravoException('Invalid or outdated credentials store. You may need to run login.')
            if not all(k in credentials for k in required):
                raise BravoException('Invalid access token entry in credentials store. You may need to run login.')
        return credstore


def write_credstore(data):
    path = os.path.join(USER_HOME, BRAVO_DIR, BRAVO_CREDSTORE)
    with open(path, 'w') as credstore:
        json.dump(data, credstore, indent = 4)
    os.chmod(path, 0600)


def login():
    try:
        bravo_response = requests.get(BRAVO_IP_API)
        if bravo_response.status_code != 200:
            raise BravoException('Error while obtaining your public IP with Bravo API authentication server.')
        bravo_response_data = bravo_response.json()
        ip = bravo_response_data['ip']
        if not credstore_exists():
            create_credstore()
            credstore = { 'active': None, 'all': {} }
        else:
            try:
                credstore = read_credstore()
                credentials = credstore['all'].get(ip, None)
                if credentials is not None and not credentials['revoked']:
                    credstore['active'] = ip
                    write_credstore(credstore)
                    print 'You are signed in.'
                    return
            except Exception:
                credstore = { 'active': None, 'all': {} }
        bravo_response = requests.get(BRAVO_AUTH_API)
        if bravo_response.status_code != 200:
            if bravo_response.status_code == 400:   
                raise BravoException(bravo_response.json().get('error', 'Failed to obtain authentication link.'))
            else:
                raise BravoException('Error while obtaining authentication link from Bravo API authentication server.')  
        bravo_response_data = bravo_response.json()
        auth_url = bravo_response_data['auth_url']
        auth_token = bravo_response_data['auth_token']
        print 'Go to the following link in your browser:\n\n{}\n'.format(bravo_response_data['auth_url'])
        print '\nContacting Bravo API for access tokens...'
        sys.stdout.flush()
        while True:
            time.sleep(5)
            bravo_response = requests.post(BRAVO_TOKEN_API, data = {'auth_token': auth_token})
            if bravo_response.status_code != 200:
                if bravo_response.status_code == 400:
                    raise BravoException(bravo_response.json().get('error', 'Failed to obtain authentication link.'))
                else:
                    raise BravoException('Error while obtaining authentication token from Bravo API authentication server.')
            bravo_response_data = bravo_response.json()
            if bravo_response_data['access_token'] is not None:
                break
        credstore['active'] = bravo_response_data['ip']
        credstore['all'][bravo_response_data['ip']] = {
            'access_token': bravo_response_data['access_token'],
            'token_type': bravo_response_data['token_type'],
            'created': int(time.time()),
            'revoked': False
        }
        write_credstore(credstore)
        print 'Done.'
        print 'You are signed in.'
    except requests.exceptions.RequestException:
        raise BravoException('Error while accessing Bravo API authentication server.')


def print_access_token():
    credstore = read_credstore()
    print credstore['all'][credstore['active']]['access_token']


def revoke():
    if not credstore_exists():
        print 'No access tokens to revoke.'
        return
    try:
        credstore = read_credstore()
        ip_revoked = list()
        for ip, credentials in credstore['all'].iteritems():
            if not credentials['revoked']:
                if not ip_revoked:
                    bravo_response = requests.get(BRAVO_REVOKE_API, params = {'access_token': credentials['access_token']})
                    if bravo_response.status_code != 200:
                        if bravo_response.status_code == 400:
                            raise BravoException(bravo_response.json().get('error', 'Failed to revoke access.'))
                        else:
                            raise BravoException('Bravo API authentication server is not accessible.')
                ip_revoked.append(ip)
                credentials['revoked'] = True
        if ip_revoked:
            write_credstore(credstore)
            print 'Access tokens have been successfully revoked.'
        else:
            print 'No access tokens to revoke.'
    except requests.exceptions.RequestException:
        raise BravoException('Error while accessing Bravo API authentication server.')


def _query_nonpaged(url, headers):
    bravo_response = requests.get(url, headers = headers)
    if bravo_response.status_code != 200:
        if bravo_response.status_code == 400:
            raise BravoException(bravo_response.json().get('error', 'Failed to query data.'))
        else: 
            bravo_response.raise_for_status()
    return bravo_response.json()


def query_meta():
    if not credstore_exists():
        print 'No access tokens found. Please login first.'
        return
    credstore = read_credstore()
    headers = { 'Authorization': 'Bearer {}'.format(credstore['all'][credstore['active']]['access_token']) }
    query_url = 'https://bravo.sph.umich.edu/freeze5/hg38/api/{}/'.format(BRAVO_API_VERSION)
    print _query_nonpaged(query_url, headers)


def parse_filter_expressions(filter):
    op_map = { '==': 'eq',  '!=': 'ne', '>': 'gt', '<': 'lt', '>=': 'gte', '<=': 'lte'}
    parsed_filter = list()
    for expression in filter.split('&'):
        tokens = re.split(r'(==|\!=|>|<|>=|<=)', expression)
        if len(tokens) != 3:
            raise BravoException('Bad filter expression.')
        parsed_filter.append('{}={}:{}'.format(tokens[0].strip(), op_map[tokens[1]], tokens[2].strip()))
    return '&'.join(parsed_filter)


def _query_paged(headers, url):
    page_no = 1
    while url:
        bravo_response = requests.get(url, headers = headers)
        if bravo_response.status_code != 200:
            if bravo_response.status_code == 400:
                raise BravoException(bravo_response.json().get('error', 'Failed to query data.'))
            else: 
                bravo_response.raise_for_status()
        bravo_response_data = bravo_response.json()
        if bravo_response_data['format'] == 'vcf' and page_no == 1:
            for line in bravo_response_data['meta']:
                yield line
            yield bravo_response_data['header']
        for item in bravo_response_data['data']:
            yield item
        url = bravo_response_data['next']
        page_no += 1


def query_region(chromosome, start, end, format, filter):
    if not credstore_exists():
        print 'No access tokens found. Please login first.'
        return
    credstore = read_credstore()
    headers = { 'Authorization': 'Bearer {}'.format(credstore['all'][credstore['active']]['access_token']) }
    query_url = 'https://bravo.sph.umich.edu/freeze5/hg38/api/{}/region?chrom={}&start={}&end={}&vcf={}'.format(BRAVO_API_VERSION, chromosome, start, end, 0 if format != 'vcf' else 1)
    if filter:
        query_url = '{}&{}'.format(query_url, parse_filter_expressions(filter))
    for line in _query_paged(headers, query_url):
        print line


def query_gene(name, format, filter):
    if not credstore_exists():
        print 'No access tokens found. Please login first.'
        return
    credstore = read_credstore()
    headers = { 'Authorization': 'Bearer {}'.format(credstore['all'][credstore['active']]['access_token']) }
    query_url = 'https://bravo.sph.umich.edu/freeze5/hg38/api/{}/gene?name={}&vcf={}'.format(BRAVO_API_VERSION, name, 0 if format != 'vcf' else 1)
    if filter:
        query_url = '{}&{}'.format(query_url, parse_filter_expressions(filter))
    for line in _query_paged(headers, query_url):
        print line


def query_variant(variant_id, chromosome, position, format):
    if not credstore_exists():
        print 'No access tokens found. Please login first.'
        return
    credstore = read_credstore()
    headers = { 'Authorization': 'Bearer {}'.format(credstore['all'][credstore['active']]['access_token']) }
    if variant_id is None:
        if chromosome is None or position is None:
            raise BravoException('Provide both "-c,--chromosome" and "-p,--position".')
        query_url = 'https://bravo.sph.umich.edu/freeze5/hg38/api/{}/variant?chrom={}&pos={}&vcf={}'.format(BRAVO_API_VERSION, chromosome, position, 0 if format != 'vcf' else 1)
    else:
        if chromosome is not None or position is not None:
            raise BravoException('"-v,--variant" is not allowed together with "-c,--chromosome" and "-p,--position".')
        query_url = 'https://bravo.sph.umich.edu/freeze5/hg38/api/{}/variant?variant_id={}&vcf={}'.format(BRAVO_API_VERSION, variant_id, 0 if format != 'vcf' else 1)
    for line in _query_paged(headers, query_url):
        print line


def load_region(chromosome, position, filter):
    if not credstore_exists():
        print 'No access tokens found. Please login first.'
        return
    credstore = read_credstore()
    headers = { 'Authorization': 'Bearer {}'.format(credstore['all'][credstore['active']]['access_token']) }
    start = position
    end = start + 8000 # approx. 1,000 variants
    region = { 'chromosome': chromosome, 'start': start, 'end': end, 'variants': dict() }
    query_url = 'https://bravo.sph.umich.edu/freeze5/hg38/api/{}/region?chrom={}&start={}&end={}&vcf=0'.format(BRAVO_API_VERSION, chromosome, start, end)
    if filter:
        query_url = '{}&{}'.format(query_url, parse_filter_expressions(filter))
    for line in _query_paged(headers, query_url):
        region['variants'][line['variant_id']] = line
    return region


def load_version():
    if not credstore_exists():
        print 'No access tokens found. Please login first.'
        return
    credstore = read_credstore()
    headers = { 'Authorization': 'Bearer {}'.format(credstore['all'][credstore['active']]['access_token']) }
    query_url = 'https://bravo.sph.umich.edu/freeze5/hg38/api/{}/'.format(BRAVO_API_VERSION)
    return _query_nonpaged(query_url, headers)['dataset']


def annotate(filter):
    if not credstore_exists():
        print 'No access tokens found. Please login first.'
        return
    region = None
    data_version = load_version()
    fileformat_line = sys.stdin.readline()
    if not fileformat_line or not fileformat_line.startswith('##fileformat=VCF'):
        return
    sys.stdout.write('{}\n'.format(fileformat_line.rstrip()))
    for in_line in sys.stdin:
        if in_line.startswith('#'):
            if in_line.startswith('##'):
                sys.stdout.write('{}\n'.format(in_line.rstrip()))
            elif in_line.startswith('#CHROM'):
                sys.stdout.write('##INFO=<ID=BRAVO_AN,Number=1,Type=Integer,Description=\"Number of Alleles in Samples with Coverage from {}\">\n'.format(data_version))
                sys.stdout.write('##INFO=<ID=BRAVO_AC,Number=A,Type=Integer,Description=\"Alternate Allele Counts in Samples with Coverage from {}\">\n'.format(data_version))
                sys.stdout.write('##INFO=<ID=BRAVO_AF,Number=A,Type=Float,Description=\"Alternate Allele Frequencies from {}\">\n'.format(data_version))
                sys.stdout.write('##INFO=<ID=BRAVO_FILTER,Number=A,Type=Float,Description=\"Filter from {}\">\n'.format(data_version))
                sys.stdout.write('{}\n'.format(in_line.rstrip()))
            continue
        in_fields = in_line.rstrip().split('\t', 8)
        chromosome = in_fields[0]
        position = int(in_fields[1])
        ref = in_fields[3]
        alt = in_fields[4] # assume bi-allelic
        info = in_fields[7]
        if region is None or region['chromosome'] != chromosome or position < region['start'] or position > region['end']:
            region = load_region(chromosome, position, filter)
        variant_id = '{}-{}-{}-{}'.format(chromosome, position, ref, alt)
        bravo_variant = region['variants'].get(variant_id, None)
        if bravo_variant is None:
            sys.stdout.write('{}\n'.format(in_line.rstrip()))
        else:
            new_info = 'BRAVO_AN={};BRAVO_AC={};BRAVO_AF={};BRAVO_FILTER={}'.format(bravo_variant['allele_num'], bravo_variant['allele_count'], bravo_variant['allele_freq'], bravo_variant['filter'])
            if info == '.':
                info = new_info
            else:
                info = '{};{}'.format(info, new_info)
            if len(in_fields) > 8:
                sys.stdout.write('{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\n'.format(chromosome, position, in_fields[2], ref, alt, in_fields[5], in_fields[6], info, in_fields[8]))
            else:
                sys.stdout.write('{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\n'.format(chromosome, position, in_fields[2], ref, alt, in_fields[5], in_fields[6], info))
 

if __name__ == '__main__':
    args = argparser.parse_args()
    try:
        if args.command == 'login':
            login()
        elif args.command == 'print-access-token':
            print_access_token()
        elif args.command == 'revoke':
            revoke()
        elif args.command == 'query-meta':
            query_meta()
        elif args.command == 'query-region':
            query_region(args.chromosome, args.start, args.end, args.format, args.filter)
        elif args.command == 'query-gene':
            query_gene(args.gene, args.format, args.filter)
        elif args.command == 'query-variant':
            query_variant(args.variant_id, args.chromosome, args.position, args.format)
        elif args.command == 'annotate':
            annotate(args.filter)
    except BravoException as e:
        print e