#!/usr/bin/env python
from __future__ import print_function

import argparse
import sys
import os

from dammit import common
from dammit import annotate
from dammit import databases
from dammit import dependencies
from dammit.tasks import print_tasks

def get_database_dir(args):
    # By default, we store databases in the home directory
    db_dir = args.database_dir
    if db_dir is None:
        db_dir = os.path.join(os.environ['HOME'], 
                              common.CONFIG['settings']['db_dir'])
    return os.path.abspath(db_dir)


def handle_databases(args):

    # TODO: Check that directory exists
    db_dir = get_database_dir(args)
    print('DB_DIR', db_dir, file=sys.stderr)

    path_dict, _ = dependencies.get_dependency_tasks()
    database_dict, tasks = databases.get_database_tasks(db_dir,
                                                        path_dict,
                                                        args.busco_group,
                                                        args.full)
    
    if args.debug:
        print_tasks(tasks)
    
    if args.install:
        databases.run_database_tasks(db_dir, tasks)
    else:
        databases.run_database_tasks(db_dir, tasks, ['list'])


def handle_dependencies(args):
    
    dep_dir = dependencies.get_dependency_dir()
    print('DEP_DIR', dep_dir, file=sys.stderr)
    try:
        os.makedirs(dep_dir)
    except OSError:
        # Already exists
        pass

    path_dict, tasks = dependencies.get_dependency_tasks()
   
    if args.debug:
        print_tasks(tasks)

    if args.install:
        dependencies.run_dependency_tasks(tasks)
    else:
        dependencies.run_dependency_tasks(tasks, ['list'])


def handle_annotate(args):

    transcriptome = args.transcriptome
    out_dir = args.output_dir
    if out_dir is None:
        out_dir = transcriptome + '.dammit'
    out_dir = os.path.abspath(out_dir)


    prog_paths, _ = dependencies.get_dependency_tasks()

    db_dir = get_database_dir(args)
    db_dict, _ = databases.get_database_tasks(db_dir, 
                                              prog_paths,
                                              args.busco_group,
                                              args.full)

    if args.debug:
        print('Transcriptome: {transcriptome}\n'\
              'Output Dir: {out_dir}\n'\
              'Paths: {prog_paths}\n'\
              'Databases: {db_dict}'.format(**locals()))

    results, tasks = annotate.get_annotate_tasks(transcriptome,
                                        prog_paths, 
                                        db_dict)

    if args.debug:
        print_tasks(tasks)

    annotate.run_annotate_tasks(transcriptome, out_dir, tasks)

def main():
    parser = argparse.ArgumentParser()
    # Debug can be used with all subcommands
    parser.add_argument('--debug', action='store_true', default=False)
    # Add subcommand parsers
    subparsers = parser.add_subparsers()
    
    '''
    `dammit databases` subcommand. This will run tasks to check for and
    install the databases. By default, only check their status; with
    `--install`, download and/or prepare them as well.
    '''
    databases_parser = subparsers.add_parser('databases')
    databases_parser.add_argument('--database-dir', default=None)
    databases_parser.add_argument('--install', action='store_true',
                                  default=False)
    databases_parser.add_argument('--full', action='store_true', default=False,
                        help='Do full annotation with uniref90')
    databases_parser.add_argument('--busco-group', default='metazoa')
    databases_parser.set_defaults(func=handle_databases)

    '''
    `dammit dependencies` subcommand. Checks for necessary executables
    on the system path and downloads them to the dammit folder if not present.
    Like the databases subcommand, takes an `install` option, which will put
    dependencies in the $HOME/.dammit/dependencies directory and give absolute
    paths to the task creators.
    '''
    dependencies_parser = subparsers.add_parser('dependencies')
    dependencies_parser.add_argument('--install', action='store_true',
                                     default=False)
    dependencies_parser.set_defaults(func=handle_dependencies)

    '''
    `dammit annotate` subcommand, which drives the primary program. The
    `--database-dir` and `--full` flags are the same as in the `databases`
    subcommand. `-o/--output-dir` specifies the directory where the results will
    be stored, and also the working directory where all annotate tasks will be
    executed. `--busco-group` specifies the BUSCO database to use; this must be
    one of BUSCOs builtins. `--n_threads` specifies the number of threads for
    each individual program to use (when applicable). `--user-databases` is an
    optional list of protein databases to run additional homology searches
    against and pull annotation information from; dammit can only take what it
    finds from these, so uninformative fasta headers will yield uninformative
    results. Finally, `transcriptome` is the transcriptome assembly to annotate.
    '''
    annotate_parser = subparsers.add_parser('annotate')
    annotate_parser.add_argument('--database-dir', default=None)
    annotate_parser.add_argument('--full', action='store_true', default=False,
                        help='Do full annotation with uniref90')
    annotate_parser.add_argument('-o', '--output-dir', default=None)
    annotate_parser.add_argument('--busco-group', default='metazoa')
    annotate_parser.add_argument('--n_threads', type=int, default=1)
    annotate_parser.add_argument('--user-databases', nargs='+')
    annotate_parser.add_argument('transcriptome')
    annotate_parser.set_defaults(func=handle_annotate)

    #args, doit_args = parser.parse_known_args()
    args = parser.parse_args()

    # Print out the infos~
    print(common.CONFIG['meta']['description'], file=sys.stderr)
    print(', '.join(common.CONFIG['meta']['authors']), common.CONFIG['meta']['date'], 
          file=sys.stderr)

    args.func(args)

if __name__ == '__main__':
    main()
