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

import argparse
import logging
import os
import sys

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

logger = logging.getLogger(__name__)

def handle_databases(args):

    common.print_header('submodule: databases', level=1)

    missing = dependencies.do_check()
    if missing:
        logger.error('Install dependencies to continue; exiting')
        return 1
    
    db_dir = databases.get_dir(args)
    logger.debug('database dir: {0}'.format(db_dir))
    try:
        os.mkdir(db_dir)
    except OSError:
        logger.debug('database dir already exists')

    database_dict, tasks, missing = databases.do_check(db_dir,
                                                       args)
    print_tasks(tasks, logger=logger)
    
    if args.install:
        if missing:
            common.print_header('Installing databases', level=2)
            databases.run_tasks(db_dir, tasks)
        else:
            common.print_header('Nothing to install', level=2)

def handle_dependencies(args):

    common.print_header('submodule: dependencies', level=1)
    logger.debug('dependencies')
   
    missing = dependencies.do_check()


def handle_annotate(args):

    common.print_header('submodule: annotate', level=1)

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

    missing_deps = dependencies.do_check()
    if missing_deps:
        logger.error('Install dependencies to continue; exiting')
        return 1
    
    db_dir = databases.get_dir(args)
    logger.debug('database dir: {0}'.format(db_dir))
    database_dict, _, missing_dbs = databases.do_check(db_dir,
                                                       args)

    if missing_dbs:
        logger.error('Install databases to continue; exiting')
        return 1

    results, tasks = annotate.get_tasks(transcriptome,
                                        database_dict,
                                        n_threads=args.n_threads,
                                        user_databases=args.user_databases)

    common.print_header('Running annotate!', level=2)
    logger.info('Transcriptome file: {0}'.format(transcriptome))
    logger.info('Output directory: {0}'.format(out_dir))
    print_tasks(tasks, logger=logger)

    annotate.run_tasks(transcriptome, out_dir, tasks)

def main():

    '''
    Build the main description string.
    '''
    meta = '{0}\n{1} {2}'.format(common.CONFIG['meta']['description'],
                                 ', '.join(common.CONFIG['meta']['authors']),
                                 common.CONFIG['meta']['date'])

    '''
    Build the main parser.
    '''
    parser = argparse.ArgumentParser(
                 description=common.add_header(meta, 0),
                 formatter_class=argparse.RawDescriptionHelpFormatter
                 )

    parser.add_argument('--debug', action='store_true', default=False)
    subparsers = parser.add_subparsers(title='dammit subcommands')

    def add_common_args(parser):
        ''' Add shared options to a parser.

        Shared options are added this way instead of to the main parser
        because I'd rather they come after the subcommand name.

        Args:
            parser (object): The parser to which arguments will be added.
        '''
        parser.add_argument('--database-dir', 
                            default=None, 
                            help='Directory to store databases. Existing'\
                                 ' databases will not be overwritten.'\
                                 ' By default, the database directory is'\
                                 ' $HOME/.dammit/databases.'
                            )

        parser.add_argument('--full', 
                            action='store_true', 
                            default=False,
                            help='Do complete run with uniref90. By default'\
                                 ' uniref90 is left out, as it is huge '\
                                 ' and homology searches take a long'\
                                 ' time.'
                            )

        parser.add_argument('--busco-group', 
                            default='metazoa',
                            choices=['metazoa', 'eukaryota', 'vertebrata',
                                     'arthropoda'],
                            help='Which BUSCO group to use. Depends on'\
                                 ' the organism being annotated.'
                            )

    '''
    Add the databases subcommand.
    '''
    desc = '''Check for databases and optionally download and prepare them 
           for use. By default, only check their status.'''
    databases_parser = subparsers.add_parser(
                           'databases',
                            description=desc,
                            help=desc
                            )

    databases_parser.add_argument('--install', 
                                  action='store_true',
                                  default=False,
                                  help='Install missing databases. Downloads'
                                       ' and preps where necessary'
                                  )

    add_common_args(databases_parser)
    databases_parser.set_defaults(func=handle_databases)

    '''
    Add the dependencies subcommand.
    '''
    desc = '''Checks for dependencies on system PATH. Unlike with the 
           databases, dependencies are not downloaded when missing and must 
           be installed by the user.'''
    dependencies_parser = subparsers.add_parser(
                              'dependencies',
                              description=desc,
                              help=desc
                              )
    dependencies_parser.set_defaults(func=handle_dependencies)

    '''
    Add the annotation subcommand.
    '''
    desc = '''The main annotation pipeline. Calculates assembly stats; 
           runs BUSCO; runs LAST against OrthoDB (and optionally uniref90), 
           HMMER against Pfam, Inferal against Rfam, and Conditional Reciprocal 
           Best-hit Blast against user databases; and aggregates all results in 
           a properly formatted GFF3 file.'''
    annotate_parser = subparsers.add_parser(
                          'annotate',
                          usage='%(prog)s <transcriptome> [OPTIONS]',
                          description=desc,
                          help=desc
                          )

    annotate_parser.add_argument('transcriptome', 
                                 help='FASTA file with the transcripts to be'\
                                      ' annotated.'
                                 )

    annotate_parser.add_argument('-o', '--output-dir', 
                                 default=None,
                                 help='Output directory. By default this will'\
                                      ' be the name of the transcriptome file'\
                                      ' with `.dammit` appended'
                                 )

    annotate_parser.add_argument('--n_threads', 
                                 type=int, 
                                 default=1,
                                 help='Number of threads to pass to programs'\
                                      ' supporting multithreading'
                                 )

    annotate_parser.add_argument('--user-databases', 
                                 nargs='+', 
                                 default=[],
                                 help='Optional additional protein databases. '\
                                      ' These will be searched with CRB-blast.'
                                 )

    add_common_args(annotate_parser)
    annotate_parser.set_defaults(func=handle_annotate)

    args = parser.parse_args()

    common.print_header(meta, 0)

    args.func(args)

if __name__ == '__main__':
    main()
