#===============================================================================
# bioleach.py
#===============================================================================

# Imports ======================================================================

from argparse import ArgumentParser
import os
from pathlib import Path
from sqlalchemy import text
import sys
from bioleach_balrog.version import __version__
from bioleach import (
    BIOLEACH_SOURCE_DB_FILE,
    BIOLEACH_PIPELINE_DB_FILE,
    BIOLEACH_DATA_DIR,
    BIOLEACH_READS_DIR,
    BIOLEACH_TRIM_ALGORITHM,
    BIOLEACH_REMOVE_CONTAM_ALGORITHM,
    BIOLEACH_CONTIGS_DIR,
    BIOLEACH_ALIGNMENTS_DIR,
    BIOLEACH_GENES_DIR
)
from bioleach import (
    dialect_agnostic_engine,
    _create_source_database,
    _create_pipeline_database,
    _export_database,
    _download_sra_reads,
    construct_read_preprocessing_func,
    _trim,
    _remove_contaminants,
    _assemble_contigs,
    mark_complete,
    _clean_up
)
from bioleach_balrog.predict_genes_balrog import _predict_genes_balrog


# Functions ====================================================================


def _predict_genes(args):
    if not Path(args.genes_dir).is_dir():
        os.mkdir(args.genes_dir)
    engine = dialect_agnostic_engine(
        dialect='mysql' if args.mysql else 'sqlite', database=args.database
    )
    with engine.connect() as conn:
        run_accessions = tuple(
            row.run_accession for row in conn.execute(text("SELECT run_accession FROM pipeline_accessions"))
        )
        for run_accession in run_accessions:
            completed = _predict_genes_balrog(
                run_accession,
                args.contigs_dir,
                args.genes_dir
            )
            if completed:
                mark_complete(run_accession, 'predicted_genes', conn)


def _data_dir(args):
    print(BIOLEACH_DATA_DIR)


def parse_arguments():
    parser = ArgumentParser(description='Bioinformatics analysis of a bioleaching microbial community')
    parser.add_argument('--version', action='version',
        version='%(prog)s {version}'.format(version=__version__))
    parser.set_defaults(func=lambda _: parser.print_help(sys.stdout))
    subparsers = parser.add_subparsers()

    # create-source-db
    parser_create_source_db = subparsers.add_parser(
        'create-source-db', help='create source database'
    )
    parser_create_source_db.set_defaults(func=_create_source_database)
    parser_create_source_db.add_argument(
        '-d', '--database', metavar='<database[.db]>', default=BIOLEACH_SOURCE_DB_FILE,
        help=f'path to sqlite database file, or name of mysql database (default: {BIOLEACH_SOURCE_DB_FILE})'
    )
    parser_create_source_db.add_argument(
        '-m', '--mysql', action='store_true',
        help='use a configured mysql database, requires that environment variables be set'
    )

    # create-pipeline-db
    parser_create_pipeline_db = subparsers.add_parser(
        'create-pipeline-db', help='create pipeline database'
    )
    parser_create_pipeline_db.set_defaults(func=_create_pipeline_database)
    parser_create_pipeline_db.add_argument(
        '-d', '--database', metavar='<database[.db]>', default=BIOLEACH_PIPELINE_DB_FILE,
        help=f'path to sqlite pipeline database file, or name of mysql pipeline database (default: {BIOLEACH_PIPELINE_DB_FILE})'
    )
    parser_create_pipeline_db.add_argument(
        '-s', '--source-db', metavar='<database[.db]>', default=BIOLEACH_SOURCE_DB_FILE,
        help=f'path to sqlite source database file, or name of mysql source database (default: {BIOLEACH_SOURCE_DB_FILE})'
    )
    parser_create_pipeline_db.add_argument(
        '-a', '--accessions', metavar='<SRRX>', nargs='*', default=[],
        help='SRA accessions to include in pipeline'
    )
    parser_create_pipeline_db.add_argument(
        '-m', '--mysql', action='store_true',
        help='use a configured mysql database, requires that environment variables be set'
    )

    # export-db
    parser_export_db = subparsers.add_parser('export-db', help='export database')
    parser_export_db.set_defaults(func=_export_database)
    parser_export_db.add_argument(
        'table', metavar='<table>', help='name of table to export'
    )
    db_select_group = parser_export_db.add_mutually_exclusive_group(required=True)
    db_select_group.add_argument(
        '-s', '--source', action='store_true', help='export the source db'
    )
    db_select_group.add_argument(
        '-p', '--pipeline', action='store_true', help='export the pipeline db'
    )
    db_select_group.add_argument(
        '-d', '--database', metavar='<database[.db]>',
        help='path to sqlite database file or mysql database name'
    )
    parser_export_db.add_argument(
        '-o', '--out', metavar='<path/to/file.{csv,tsv}', default=sys.stdout,
        help='path to output file'
    )
    parser_export_db.add_argument(
        '-m', '--mysql', action='store_true',
        help='use a configured mysql database, requires that environment variables be set'
    )

    # download
    parser_download = subparsers.add_parser('download', help='download SRA reads')
    parser_download.set_defaults(func=_download_sra_reads)
    parser_download.add_argument(
        '--pipeline-db', metavar='<database[.db]>', default=BIOLEACH_PIPELINE_DB_FILE,
        help=f'path to sqlite pipeline database file, or mysql pipeline database name (default: {BIOLEACH_PIPELINE_DB_FILE})'
    )
    parser_download.add_argument(
        '-m', '--mysql', action='store_true',
        help='use a configured mysql database, requires that environment variables be set'
    )

    # trim and preprocess
    parser_trim = subparsers.add_parser(
        'trim', help='trim Illumina adapters and low-quality bases'
    )
    parser_trim.set_defaults(func=construct_read_preprocessing_func(_trim, 'trimmed'))
    parser_trim.add_argument(
        '--database', metavar='<database[.db]>', default=BIOLEACH_PIPELINE_DB_FILE,
        help=f'path to sqlite pipeline database file, or mysql pipeline database name (default: {BIOLEACH_PIPELINE_DB_FILE})'
    )
    parser_trim.add_argument(
        '--output-dir', metavar='<output_dir/>', default=BIOLEACH_READS_DIR,
        help=f'destination directory for trimmed reads (default: {BIOLEACH_READS_DIR})'
    )
    parser_trim.add_argument(
        '--algorithm', choices=('fastp', 'bbtools', 'trim-galore', 'trimmomatic'), default=BIOLEACH_TRIM_ALGORITHM,
        help=f'choose algorithm for adapter and quality trimming (default: {BIOLEACH_TRIM_ALGORITHM})'
    )
    parser_trim.add_argument(
        '--processes', type=int, default=1,
        help='number of processes for adapter and quality trimming (default: 1)'
    )
    parser_trim.add_argument(
        '--parallel', action='store_true',
        help='trim multiple sequencing datasets in parallel'
    )
    parser_trim.add_argument(
        '-m', '--mysql', action='store_true',
        help='use a configured mysql database, requires that environment variables be set'
    )
    
    # remove contaminants
    parser_remove_contaminants = subparsers.add_parser(
        'remove-contam', help='remove contaminant reads'
    )
    parser_remove_contaminants.set_defaults(
        func=construct_read_preprocessing_func(
            _remove_contaminants,
            'removed_contaminants',
            prev_fastq_suffix='_{0}_val_{0}.fq'
        )
    )
    parser_remove_contaminants.add_argument(
        '--database', metavar='<database[.db]>', default=BIOLEACH_PIPELINE_DB_FILE,
        help=f'path to sqlite pipeline database file, or mysql pipeline database name (default: {BIOLEACH_PIPELINE_DB_FILE})'
    )
    parser_remove_contaminants.add_argument(
        '--output-dir', metavar='<output_dir/>', default=BIOLEACH_READS_DIR,
        help=f'destination directory for decontaminated reads (default: {BIOLEACH_READS_DIR})'
    )
    parser_remove_contaminants.add_argument(
        '--algorithm', choices=('kraken2', 'bbtools', 'bowtie2', 'bwa', 'minimap2'),
        default=BIOLEACH_REMOVE_CONTAM_ALGORITHM,
        help=f'choose algorithm for contaminant removal (default: {BIOLEACH_REMOVE_CONTAM_ALGORITHM})'
    )
    parser_remove_contaminants.add_argument(
        '--processes', type=int, default=1,
        help='number of processes for contaminant removal (default: 1)'
    )
    parser_remove_contaminants.add_argument(
        '-m', '--mysql', action='store_true',
        help='use a configured mysql database, requires that environment variables be set'
    )

    # assemble contigs
    parser_assemble_contigs = subparsers.add_parser(
        'assemble',
        help='assemble metagenome contigs'
    )
    parser_assemble_contigs.set_defaults(func=_assemble_contigs)
    parser_assemble_contigs.add_argument(
        '--database', metavar='<database[.db]>', default=BIOLEACH_PIPELINE_DB_FILE,
        help=f'path to sqlite pipeline database file, or mysql pipeline database name (default: {BIOLEACH_PIPELINE_DB_FILE})'
    )
    parser_assemble_contigs.add_argument(
        '--contigs-dir', metavar='<contigs_dir/>', default=BIOLEACH_CONTIGS_DIR,
        help=f'destination directory for assembled contigs (default: {BIOLEACH_CONTIGS_DIR})'
    )
    parser_assemble_contigs.add_argument(
        '--alignments-dir', metavar='<alignments_dir/>', default=BIOLEACH_ALIGNMENTS_DIR,
        help=f'destination directory for BAM files (default: {BIOLEACH_ALIGNMENTS_DIR})'
    )
    parser_assemble_contigs.add_argument(
        '--processes', type=int, default=1,
        help='number of processes for contig assembly (default: 1)'
    )
    parser_assemble_contigs.add_argument(
        '-m', '--mysql', action='store_true',
        help='use a configured mysql database, requires that environment variables be set'
    )

    # assemble contigs
    parser_predict_genes = subparsers.add_parser(
        'predict',
        help='predict genes'
    )
    parser_predict_genes.set_defaults(func=_predict_genes)
    parser_predict_genes.add_argument(
        '--database', metavar='<database[.db]>', default=BIOLEACH_PIPELINE_DB_FILE,
        help=f'path to sqlite pipeline database file, or mysql pipeline database name (default: {BIOLEACH_PIPELINE_DB_FILE})'
    )
    parser_predict_genes.add_argument(
        '--contigs-dir', metavar='<contigs_dir/>', default=BIOLEACH_CONTIGS_DIR,
        help=f'destination directory for assembled contigs (default: {BIOLEACH_CONTIGS_DIR})'
    )
    parser_predict_genes.add_argument(
        '--genes-dir', metavar='<genes_dir/>', default=BIOLEACH_GENES_DIR,
        help=f'destination directory for BAM files (default: {BIOLEACH_GENES_DIR})'
    )
    parser_predict_genes.add_argument(
        '-m', '--mysql', action='store_true',
        help='use a configured mysql database, requires that environment variables be set'
    )

    # clean up
    parser_clean = subparsers.add_parser('clean', help='clean up SRA reads')
    parser_clean.set_defaults(func=_clean_up)
    parser_clean.add_argument(
        '--directory', metavar='<directory/>',
        default=BIOLEACH_READS_DIR,
        help=f'directory containing downloaded SRA reads (default: {BIOLEACH_READS_DIR})'
    )
    parser_dialect_group = parser_clean.add_mutually_exclusive_group()
    parser_dialect_group.add_argument(
        '-s', '--sqlite', action='store_true',
        help='clean up a configured sqlite database'
    )
    parser_dialect_group.add_argument(
        '-m', '--mysql', action='store_true',
        help='clean up a configured mysql database, requires that environment variables be set'
    )
    parser_clean.add_argument(
        '--source-db', metavar='<database[.db]>', default=BIOLEACH_SOURCE_DB_FILE,
        help=f'path to sqlite pipeline database file, or mysql pipeline database name (default: {BIOLEACH_SOURCE_DB_FILE})'
    )
    parser_clean.add_argument(
        '--pipeline-db', metavar='<database[.db]>', default=BIOLEACH_PIPELINE_DB_FILE,
        help=f'path to sqlite pipeline database file, or mysql pipeline database name (default: {BIOLEACH_PIPELINE_DB_FILE})'
    )

    # data directory
    parser_data_dir = subparsers.add_parser(
        'data-dir', help='show the currently configured data directory'
    )
    parser_data_dir.set_defaults(func=_data_dir)
    return parser.parse_args()


def main():
    args = parse_arguments()
    args.func(args)


# Execute ======================================================================

if __name__ == '__main__':
    main()
