#!/usr/bin/env python
'''
Basepair CLI, wraps the Python API

[2015] - [2019] Basepair INC
All Rights Reserved.

NOTICE: All information contained herein is, and remains
the property of Basepair INC and its suppliers,
if any. The intellectual and technical concepts contained
herein are proprietary to Basepair INC
and its suppliers and may be covered by U.S. and Foreign Patents,
patents in process, and are protected by trade secret or copyright law.
Dissemination of this information or reproduction of this material
is strictly forbidden unless prior written permission is obtained
from Basepair INC.
'''

from __future__ import print_function

import argparse
import json
import sys
import os

# App imports
from basepair.helpers import eprint

sys.path.insert(0, '/home/ec2-user/basepair')
import basepair


def yes_or_no(question):
  '''
  Ask a yes/no question via raw_input() and return their answer.

  Parameters
  ----------
  question - str
      Is a string that is presented to the user.

  default - str
      is the presumed answer if the user just hits <Enter>.
      It must be "yes" (the default), "no" or None (meaning
      an answer is required of the user).

  '''
  valid = {
    'yes': True,
    'y': True,
    'ye': True,
    'no': False,
    'n': False
  }

  prompt = ' [y/n] '

  while True:

    # get input function for py 2 and 3
    get_input = input
    if sys.version_info[:2] <= (2, 7):
      get_input = raw_input

    # get input from user
    sys.stdout.write(question + prompt)
    choice = get_input().lower()

    # check answer
    if choice in valid:
      return valid[choice]
    sys.stdout.write('Please respond with \'yes\' or \'no\' (or \'y\' or \'n\').\n')


def main():
  '''Main method'''
  args = read_args()

  if args.config:
    conf = json.load(open(args.config))
  else:
    conf = None

  # will override default and use these as owner for all actions
  # if args.username:
  #     conf['api']['username'] = args.username
  # if args.api_key:
  #     conf['api']['api_key'] = args.api_key

  bp = basepair.connect(
    conf=conf,
    scratch=args.scratch,
    use_cache=args.use_cache,
    user_cache_for_host_conf=args.keep_cloud_service_conf,
    verbose=args.verbose
  )

  if args.action == 'create':
    if args.create_type == 'project':
      create_project(bp, args)

    if args.create_type == 'sample':

      # if payload username or api key specified, make sure both are present
      if args.payload_username is not None and args.payload_api_key is None:
        eprint('specify parameter --payload-api-key!')
        sys.exit(1)
      elif args.payload_api_key is not None and args.payload_username is None:
        eprint('specify parameter --payload-username!')
        sys.exit(1)
      elif args.payload_username is not None and args.payload_username is not None:
        bp.payload = {
          'username': args.payload_username,
          'api_key': args.payload_api_key,
        }
      create_sample(bp, args)
    if args.create_type == 'analysis':
      create_analysis(bp, args)
  elif args.action == 'updateProject':
    update_project(bp, args)
  elif args.action == 'updateSample':
    update_sample(bp, args)
  elif args.action == 'updateAnalysis':
    update_analysis(bp, args.analysis, args.key, args.val)
  elif args.action == 'download':
    if args.download_type == 'analysis':
      download_analysis(bp, args)

    elif args.download_type == 'file':
      download_file(bp, args)

    elif args.download_type == 'log':
      download_log(bp, args)

    elif args.download_type == 'sample':
      download_sample(bp, args)

  elif args.action == 'delete':
    delete(bp, args.type, args.uid)
  elif args.action == 'list':

    if args.list_type == 'analysis':
      _list_item(bp, data_type='analysis', uids=args.uid, is_json=args.json)
    if args.list_type == 'analyses':
      bp.print_data(data_type='analyses', is_json=args.json, project=args.project)
    elif args.list_type == 'genomes':
      bp.print_data(data_type='genomes', is_json=args.json)
    elif args.list_type == 'sample':
      _list_item(bp, data_type='sample', uids=args.uid, is_json=args.json)
    elif args.list_type == 'samples':
      bp.print_data(data_type='samples', is_json=args.json, project=args.project)
    elif args.list_type == 'workflows':
      bp.print_data(data_type='workflows', is_json=args.json)
    elif args.list_type == 'genome':
      _list_item(bp, data_type='genome', uids=args.uid, is_json=args.json)


def create_project(bp, args):
  '''Create project'''
  data = {'name': args.name}
  bp.create_project(data=data)


def create_sample(bp, args):
  '''Create sample'''
  data = {
    'datatype': args.datatype,
    'default_workflow': int(args.workflow) if args.workflow else None,
    'filepaths1': args.file1,
    'filepaths2': args.file2,
    'genome': args.genome,
    'name': args.name,
    'platform': args.platform,
    'projects': int(args.project) if args.project else None,
  }

  if args.key and args.val:
    for key, val in zip(args.key, args.val):
      data[key] = val

  bp.create_sample(data, upload=True, source='cli')


def create_analysis(bp, args):
  '''Create and submit an analysis'''
  params = {'node': {}}

  if not args.workflow:
    eprint('ERROR: Workflow required.')
    return

  if not args.sample:
    eprint('ERROR: Minimum one sample required.')
    return

  if args.params:
    for param in args.params:
      node_id, arg, val = param.split(':')
      if node_id not in params['node']:
        params['node'][node_id] = {}
      params['node'][node_id][arg] = val
  else:
    eprint('You specified no parameters, submitting with default ones...')

  bp.create_analysis(
    control_ids=args.control or [],
    ignore_validation_warnings=args.ignore_warning,
    params=params,
    project_id=args.project,
    sample_ids=args.sample,
    workflow_id=args.workflow,
  )


def download_analysis(bp, args):
  '''Download analysis'''
  if not args.uid:
    eprint('ERROR: Minimum one analysis uid required.')
    return

  # download a file from an analysis by tags
  for uid in args.uid:
    bp.download_analysis(uid, outdir=args.outdir, tagkind=args.tagkind, tags=args.tags)


def download_file(bp, args):
  '''Download file by uid'''
  if not args.uid:
    eprint('ERROR: Minimum one file uid required.')
    return

  for uid in args.uid:
    file_i = bp.get_file(uid)
    bp.download_file(file_i['path'], dirname=args.outdir)


def download_log(bp, args):
  '''Download analysis log'''
  if not args.uid:
    eprint('ERROR: Minimum one analysis uid required.')
    return

  for uid in args.uid:
    info, res = bp.get_analysis(uid)  # check analysis id is valid
    if info is None:
      eprint('{} is not a valid analysis id!'.format(uid))
      continue

    if args.outdir:
      bp.get_log(uid, args.outdir)
    else:
      bp.get_log(uid)


def download_sample(bp, args):
  '''Download sample'''
  if not args.uid:
    eprint('ERROR: Minimum one sample uid required.')
    return

  for uid in args.uid:
    sample = bp.get_sample(uid, add_analysis=True)  # check sample id is valid
    if sample is None:
      eprint('{} is not a valid sample id!'.format(uid))
      continue

    # if tags provided, download file by tags
    if args.tags:
      bp.get_file_by_tags(sample, tags=args.tags, kind=args.tagkind, dirname=args.outdir)
    else:
      bp.download_raw_files(sample, args.outdir)


def update_info(bp, kind=None, uid=None, keys=None, vals=None, data={}):
  '''Update object info'''
  if keys and vals:
    for key, val in zip(keys, vals):
      if kind == 'sample' and key in ['adapter', 'amplicon', 'barcode', 'regions', 'stranded']:
        data['info'] = data.get('info', {})
        data['info'][key] = val  # set sample info field
      else:
        data[key] = val

  if kind == 'sample':
    res = bp.update_sample(uid, data)
  elif kind == 'analysis':
    res = bp.update_analysis(uid, data)
  if res.get('error'):
    eprint('error: {}'.format(res.get('msg')))


def update_project(bp, args):
  '''Update project'''
  if not args.project:
    eprint('ERROR: Minimum one project required.')
    return

  data = {}
  if args.name:
    data = {'name': args.name}

  params = {}
  if args.emails and args.perm:
    params = {
      'params': json.dumps({
        'permission_data': {
          'emails': args.emails,
          'perm': args.perm,
        }
      })
    }

  if not data and not params:
    eprint('WARNING: No data to update.')
    return

  for project_id in args.project:
    res = bp.update_project(project_id, data=data, params=params)

    if res.get('error'):
      eprint('error: {}'.format(res.get('msg')))


def update_sample(bp, args):
  '''Update sample'''
  if not args.sample:
    eprint('ERROR: Sample required.')
    return

  data = {}
  if args.name:
    data['name'] = args.name

  if args.genome:
    data['genome'] = args.genome

  if args.datatype:
    data['datatype'] = args.datatype

  update_info(bp, kind='sample', uid=args.sample, keys=args.key, vals=args.val, data=data)


def update_analysis(bp, analysis_id, keys, vals):
  '''Update analysis'''
  if not analysis_id:
    eprint('Error: Analysis required.')
    return
  update_info(bp, kind='analysis', uid=analysis_id, keys=keys, vals=vals)


def delete(bp, data_type, uid):
  '''Delete object'''
  if not uid:
    eprint('Please add one or more uid')
    return

  for uid in uid:
    answer = yes_or_no('Are you sure you want to delete {}?'.format(uid))
    if answer:
      if data_type == 'sample':
        bp.delete_sample(uid)
      elif data_type == 'analysis':
        bp.delete_analysis(uid)


def add_common_args(parser):
  '''Add common args'''
  parser.add_argument('-c', '--config', help='API config info')
  parser.add_argument('--quiet', action='store_true')
  parser.add_argument('--keep-cloud-service-conf', action='store_true')
  parser.add_argument('--use-cache', action='store_true')
  parser.add_argument('--scratch', help='Scratch dir for files')
  parser.add_argument('--verbose', action='store_true')
  parser.set_defaults(scratch='.')
  return parser


def add_payload_args(parser):
  '''Add payload args'''
  parser.add_argument(
    '--payload-username',
    help='Replace payload api key. Must also specify --payload-api-key'
  )
  parser.add_argument(
    '--payload-api-key',
    help='Replace payload api key. Must also specify --payload-username'
  )
  return parser


def add_uid_parser(parser):
  '''Add uid parser'''
  parser.add_argument(
    '-u',
    '--uid',
    nargs='+',
    default=None,
    help='The unique analysis id(s) to download'
  )
  return parser


def add_json_parser(parser):
  '''Add json parser'''
  parser.add_argument(
    '--json',
    action='store_true',
    help='(Optional) Print the data in JSON format (this shows all data associated with each item).'
  )
  return parser


def add_outdir_parser(parser):
  '''Add outdir parser'''
  parser.add_argument(
    '-o',
    '--outdir',
    required=False,
    default='.',
    help='Base directory to save files to (default \'.\').'
  )
  return parser


def add_tags_parser(parser):
  '''Add tags parser'''
  parser.add_argument(
    '--tags',
    nargs='+',
    action='append',
    default=None
  )
  parser.add_argument(
    '--tagkind',
    choices=['diff', 'exact', 'subset'],
    default='exact',
    help='''
             Tag filtering strategy. Options:\n
             (1) exact (default): only get files with exactly these tags.
             (2) diff: exclude files with this tag.
             (3) subset: any files with one or more of these tags.
             '''
  )
  return parser


def read_args():
  '''Read args'''
  parser = argparse.ArgumentParser(
    description='Basepair CLI, API version {}'.format(basepair.__version__),
    formatter_class=argparse.RawDescriptionHelpFormatter
  )

  actions_p = parser.add_subparsers(
    help='Basepair actions.',
    dest='action'
  )
  actions_p.required = True

  # create parser
  create_p = actions_p.add_parser(
    'create',
    help='Create samples, analyses, etc.'
  )
  create_sp = create_p.add_subparsers(
    help='Type of item to download',
    dest='create_type'
  )

  # create project parser
  create_project_p = create_sp.add_parser(
    'project',
    help='Add a project to your account on Basepair.'
  )
  create_project_p = add_common_args(create_project_p)
  create_project_p.add_argument('--name')
  create_project_p = add_payload_args(create_project_p)

  # create sample parser
  create_sample_p = create_sp.add_parser(
    'sample',
    help='Add a sample to your account on Basepair.'
  )
  create_sample_p = add_common_args(create_sample_p)
  create_sample_p.add_argument('--name')
  create_sample_p.add_argument('--genome')
  create_sample_p.add_argument('--platform')
  create_sample_p.add_argument(
    '--datatype',
    choices=[
      'atac-seq', 'chip-seq', 'crispr', 'cutnrun', 'cutntag', 'dna-seq', 'other',
      'panel', 'rna-seq', 'scrna-seq', 'small-rna-seq', 'snap-chip', 'wes', 'wgs'
    ],
    default='rna-seq'
  )
  create_sample_p.add_argument('--file1', nargs='+')
  create_sample_p.add_argument('--file2', nargs='+')
  create_sample_p.add_argument('-w', '--workflow', help='Workflow id')
  create_sample_p.add_argument('-p', '--project', help='Project id')
  create_sample_p.add_argument('--key', action='append')
  create_sample_p.add_argument('--val', action='append')
  create_sample_p = add_payload_args(create_sample_p)

  # create analysis parser
  create_analysis_p = create_sp.add_parser(
    'analysis',
    help='Create and run an analysis.'
  )
  create_analysis_p = add_common_args(create_analysis_p)
  create_analysis_p.add_argument('-p', '--project', help='Project id')
  create_analysis_p.add_argument(
    '-w', '--workflow', help='Workflow id'
  )
  create_analysis_p.add_argument(
    '--sample', nargs='+', help='Sample id'
  )
  create_analysis_p.add_argument(
    '--control', nargs='+', help='Control id'
  )
  create_analysis_p.add_argument('--params', nargs='+')
  create_analysis_p.add_argument(
    '-i', '--ignore_warning',
    action='store_true',
    default=False,
    help='Ignore validation warnings',
  )

  # update project parser
  update_project_parser = actions_p.add_parser(
    'updateProject',
    help='Update information associated with a project.'
  )
  update_project_parser = add_common_args(update_project_parser)
  update_project_parser.add_argument(
    '-p', '--project', nargs='+', help='project id'
  )
  update_project_parser.add_argument('--emails', default=[], nargs='+')
  update_project_parser.add_argument('--perm', choices=['admin', 'edit', 'view'], default='view')
  update_project_parser.add_argument('--name')

  # update sample parser
  update_sample_parser = actions_p.add_parser(
    'updateSample',
    help='Update information associated with a sample.'
  )
  update_sample_parser = add_common_args(update_sample_parser)
  update_sample_parser.add_argument(
    '-s', '--sample', help='Sample id'
  )
  update_sample_parser.add_argument('--name')
  update_sample_parser.add_argument('--genome')
  update_sample_parser.add_argument(
    '--datatype',
    choices=[
      'atac-seq', 'chip-seq', 'crispr', 'cutnrun', 'cutntag', 'dna-seq', 'other',
      'panel', 'rna-seq', 'scrna-seq', 'small-rna-seq', 'snap-chip', 'wes', 'wgs'
    ],
  )
  update_sample_parser.add_argument('--key', action='append')
  update_sample_parser.add_argument('--val', action='append')

  # update analysis parser
  update_analysis_parser = actions_p.add_parser(
    'updateAnalysis',
    help='Update information associated with an analysis.'
  )
  update_analysis_parser = add_common_args(update_analysis_parser)
  update_analysis_parser.add_argument('-a', '--analysis', help='Analysis id')
  update_analysis_parser.add_argument('--key', action='append')
  update_analysis_parser.add_argument('--val', action='append')

  # download parsers
  download_p = actions_p.add_parser('download', help='Download items.')

  download_sp = download_p.add_subparsers(
    help='Type of item to download',
    dest='download_type'
  )

  download_analysis_p = download_sp.add_parser(
    'analysis',
    help='Download files for one or more analyses. Can filter by file tags.'
  )
  download_analysis_p = add_uid_parser(download_analysis_p)
  download_analysis_p = add_tags_parser(download_analysis_p)
  download_analysis_p = add_outdir_parser(download_analysis_p)
  download_analysis_p = add_common_args(download_analysis_p)

  download_file_p = download_sp.add_parser(
    'file',
    help='Download one or more files by uid.'
  )
  download_file_p = add_uid_parser(download_file_p)
  download_file_p = add_outdir_parser(download_file_p)
  download_file_p = add_common_args(download_file_p)

  download_log_p = download_sp.add_parser(
    'log',
    help='Download analysis logs from one or more analyses.'
  )
  download_log_p = add_uid_parser(download_log_p)
  download_log_p = add_outdir_parser(download_log_p)
  download_log_p = add_common_args(download_log_p)

  download_sample_p = download_sp.add_parser(
    'sample',
    help='Download raw data or analysis files by tags from one or more samples.'
  )
  download_sample_p = add_uid_parser(download_sample_p)
  download_sample_p = add_tags_parser(download_sample_p)
  download_sample_p = add_outdir_parser(download_sample_p)
  download_sample_p = add_common_args(download_sample_p)

  # delete parser
  delete_parser = actions_p.add_parser(
    'delete',
    help='Delete a sample or analysis from your Basepair account.'
  )
  delete_parser = add_common_args(delete_parser)
  delete_parser.add_argument(
    '-t',
    '--type',
    choices=['analysis', 'sample'],
    help='Type of data to delete. Options: sample, analysis',
    required=True,
  )
  delete_parser.add_argument(
    '-u',
    '--uid',
    default=None,
    help=' The unique sample or analysis id(s)',
    nargs='+',
  )

  # list parser
  list_p = actions_p.add_parser('list', help='List of items.')

  list_sp = list_p.add_subparsers(
    dest='list_type',
    help='Type of item to list',
  )

  list_analysis_p = list_sp.add_parser(
    'analysis',
    help='List files and other info for one or more analyses.'
  )
  list_analysis_p = add_common_args(list_analysis_p)
  list_analysis_p = add_uid_parser(list_analysis_p)
  list_analysis_p = add_json_parser(list_analysis_p)

  list_analyses_p = list_sp.add_parser(
    'analyses',
    help='List basic info about your analyses'
  )
  list_analyses_p.add_argument(
    '-p',
    '--project',
    help='List analyses of the project'
  )
  list_analyses_p = add_common_args(list_analyses_p)
  list_analyses_p = add_json_parser(list_analyses_p)

  list_samples_p = list_sp.add_parser(
    'samples',
    help='List basic info about all your samples.'
  )
  list_samples_p.add_argument(
    '-p',
    '--project',
    help='List samples of the project'
  )
  list_samples_p = add_common_args(list_samples_p)
  list_samples_p = add_json_parser(list_samples_p)

  list_sample_p = list_sp.add_parser(
    'sample',
    help='List detail info about one or more samples.'
  )
  list_sample_p = add_common_args(list_sample_p)
  list_sample_p = add_uid_parser(list_sample_p)
  list_sample_p = add_json_parser(list_sample_p)

  list_genomes_p = list_sp.add_parser(
    'genomes',
    help='List available genomes.'
  )
  list_genomes_p = add_common_args(list_genomes_p)
  list_genomes_p = add_json_parser(list_genomes_p)

  list_genome_p = list_sp.add_parser(
    'genome',
    help='List detail info about one or more genomes.'
  )
  list_genome_p = add_common_args(list_genome_p)
  list_genome_p = add_uid_parser(list_genome_p)
  list_genome_p = add_json_parser(list_genome_p)

  list_workflows_p = list_sp.add_parser(
    'workflows',
    help='List available workflows.'
  )
  list_workflows_p = add_common_args(list_workflows_p)
  list_workflows_p = add_json_parser(list_workflows_p)

  parser.set_defaults(params=[],)
  args = parser.parse_args()

  if not args.config:
    if 'BP_CONFIG_FILE' not in os.environ:
      eprint('Please either use the -c or --config param or set the environment variable BP_CONFIG_FILE!')
      sys.exit(1)
    else:
      eprint('Using config file {}'.format(os.environ['BP_CONFIG_FILE']))
  else:
    eprint('Using config file', args.config)

  # if neither set, be verbose
  if not args.verbose and not args.quiet:
    args.verbose = True

  if hasattr(args, 'key') and hasattr(args, 'val'):
    msg = None
    if args.key and not args.val:
      msg = 'val required for key'
    elif args.val and not args.key:
      msg = 'key required for val'
    elif args.key and args.val and len(args.key) != len(args.val):
      msg = 'number of key and val are not equal'

    if msg:  # stop execution if key/val error
      eprint(msg)
      sys.exit(1)

  return args

def _list_item(bp, data_type=None, uids=None, is_json=False):
  '''List single item of a type'''
  if not data_type:
    eprint('Invalid data type.')
    return

  list_type = ['analysis', 'genome', 'sample']
  if data_type not in list_type:
    eprint('List data type - {} currently not supported.'.format(data_type))
    return

  if not uids:
    eprint('At least one uid required.')
    return

  for uid in uids:
    bp.print_data(data_type=data_type, uid=uid, is_json=is_json)

if __name__ == '__main__':
  main()
