#!/usr/bin/env python
import os
import sys
from datetime import datetime
from glob import glob
from os.path import isfile, join, dirname, abspath
import shutil
import click
import subprocess

from ngs_utils.call_process import run_simple
from ngs_utils.file_utils import verify_file, safe_mkdir, verify_dir
from ngs_utils import logger
from ngs_utils.logger import warn, info
from ngs_utils.utils import set_locale; set_locale()
from ngs_utils import snakemake_utils
from umccrise import package_path


@click.command()
@click.argument('bcbio_project', type=click.Path(exists=True))
@click.argument('target_rule', nargs=-1)
@click.option('-o', 'output_dir', type=click.Path(), help='Output directory (default is "umccrise")')

@click.option('-s', '--sample', 'sample', help='Comma-separated list of samples or batches to process')
@click.option('-b', '--batch', 'batch', help='Comma-separated list of samples or batches to process')
@click.option('-e', '--exclude', 'exclude', help='Comma-separated list of samples or batches to ignore')

@click.option('-j', '--jobs', 'jobs', default=1, help='Maximum number of cores to use at single time (works both for '
              'local and cluster runs)')
@click.option('-c', '--cluster-auto', 'cluster', is_flag=True, help='Submit jobs to cluster')
@click.option('--cluster', '--cluster-cmd', 'cluster_cmd', help='Deprecated. Use --cluster-auto instead')

@click.option('--unlock', is_flag=True, help='Propagaded to snakemake')
@click.option('--rerun-incomplete', is_flag=True, help='Propagaded to snakemake')
@click.option('--restart-times', default=3, help='Propagaded to snakemake. Default is 3.')

@click.option('--pon', '--panel-of-normals', '--pon-dir', 'pon_dir', help='Path to panel of normals folder '
              '(see https://github.com/umccr/vcf_stuff#building-the-panel on how to prepare it)')
@click.option('--bcbio-genomes', help='Path to bcbio-nextgen reference data (e.g. /bcbio/genomes; '
              'Hsapiens/{genome}/seq/{genome}.fa(.fai) and Hsapiens/{genome}/validation/giab-NA12878/truth_regions.bed '
              'are used)')
@click.option('--ref-fasta', help='Path to reference fasta (e.g. /genomes/hg19.fa); .fai file should exist')
@click.option('--truth-regions', help='Path to regions to subset calls, e.g. giab truth regions (e.g. '
                                      'validation/giab-NA12878/truth_regions.bed)')
@click.option('--docker', '--docker-wrapper-mode', 'docker_wrapper_mode', is_flag=True,
              help='Pull umccr/umccrise docker image and use it instead of the main codebase')
@click.option('--pcgr-data', help='Path to the PCGR reference data folder (for dockerized version only)')
@click.option('--pcgr', help='Path to PCGR installation dir')
@click.option('--no-s3', is_flag=True, help='Do not attempt to upload results to s3')

def main(bcbio_project, target_rule=list(), output_dir=None, sample=None, batch=None, exclude=None,
         jobs=None, cluster=False, cluster_cmd=None, unlock=False, rerun_incomplete=False, restart_times=None,
         pon_dir=None, bcbio_genomes=None, ref_fasta=None, truth_regions=None,
         docker_wrapper_mode=False, pcgr_data=None, pcgr=None, no_s3=None):
    """
Umccrise (post-process) a bcbio project.\n
BCBIO_PROJECT: path to a bcbio run (final or "datestamp" directory)\n
TARGET_RULE: optional list of rules, e.g.: pcgr coverage structural small_variants rmd igv
"""

    output_dir = output_dir or 'umccrised'
    output_dir = safe_mkdir(abspath(output_dir))
    log_dir = safe_mkdir(join(output_dir, 'log'))
    logger.init(log_fpath_=join(log_dir, 'umccrise.log'), save_previous=True)

    conf = dict()
    docker_mounts = dict()

    #######################
    #### Setting paths ####
    #######################

    if docker_wrapper_mode:
        docker_mounts[output_dir] = '/output_dir'

    conf['bcbio_project'] = verify_dir(bcbio_project, is_critical=True)
    if docker_wrapper_mode:
        conf['bcbio_project'] = docker_mounts[bcbio_project] = '/bcbio_project'

    if pon_dir:
        verify_file(join(pon_dir, 'panel_of_normals.snps.vcf.gz'), is_critical=True, description='Panel of normals SNPs file in user provided folder')
        verify_file(join(pon_dir, 'panel_of_normals.snps.vcf.gz.tbi'), is_critical=True, description='Please index panel of normal files with tabix')
        verify_file(join(pon_dir, 'panel_of_normals.indels.vcf.gz'), is_critical=True, description='Panel of normals indels file in user provided folder')
        verify_file(join(pon_dir, 'panel_of_normals.indels.vcf.gz.tbi'), is_critical=True, description='Please index panel of normal files with tabix')
        conf['panel_of_normals'] = verify_dir(pon_dir, is_critical=True)
        if docker_wrapper_mode:
            conf['panel_of_normals'] = docker_mounts[pon_dir] = '/panel_of_normals'

    if bcbio_genomes:
        conf['bcbio_genomes'] = verify_dir(bcbio_genomes, is_critical=True)
        if docker_wrapper_mode:
            conf['bcbio_genomes'] = docker_mounts[bcbio_genomes] = '/bcbio_genomes'

    if ref_fasta:
        verify_file(ref_fasta + '.fai', is_critical=True)
        conf['ref_fasta'] = verify_file(ref_fasta, is_critical=True)
        if docker_wrapper_mode:
            conf['ref_fasta'] = docker_mounts[ref_fasta] = '/ref.fa'

    if truth_regions:
        conf['truth_regions'] = verify_file(truth_regions, is_critical=True)
        if docker_wrapper_mode:
            conf['truth_regions'] = docker_mounts[truth_regions] = '/truth_regions.bed'

    if pcgr_data:
        assert docker_wrapper_mode, '--pcgr-data is only for dockerized more (set the --docker flag)'
        pcgr_data = verify_dir(pcgr_data, is_critical=True, description='PCGR data directory')
        logger.info(f'Running PCGR in a dockerized mode, assuming that {pcgr_data} if a path to PCGR data directory')
        docker_mounts[pcgr_data] = '/pcgr/data'
        conf['pcgr'] = '/pcgr'
    if pcgr:
        conf['pcgr'] = verify_dir(pcgr, is_critical=True, description='PCGR installation directory')
        verify_dir(join(pcgr, 'data'), is_critical=True, description='PCGR data directory')

    if docker_wrapper_mode:
        if not bcbio_genomes and not ref_fasta:
            logger.critical('Reference fasta path is required when running in a dockerized mode. Please provide it with --ref-fasta or path to bcbio folder with --bcbio-genomes')
        if not bcbio_genomes and not truth_regions:
            logger.critical('Truth region path required when running in a dockerized mode. Please provide it with --truth-regions or path to bcbio folder with --bcbio-genomes')

    #####################################
    #### Setting non-path parameters ####
    #####################################

    if sample or batch:
        conf['sample'] = sample or batch
    if exclude:
        conf['exclude'] = exclude
    target_rule = list(target_rule)
    if 'pcgr_download' in target_rule:
        conf['pcgr_download'] = 'yes'

    if not no_s3:
        if os.environ.get('AWS_PROFILE') != 'umccr':
            warn('To upload minibams to s3://umccr-igv, please set AWS_PROFILE=umccr')
            info()
            no_s3 = True
    conf['no_s3'] = 'yes' if no_s3 else 'no'

    #########################
    #### Setting cluster ####
    #########################

    cluster_param = ''
    cluster_log_dir = ''
    if (cluster or cluster_cmd) and not docker_wrapper_mode:
        if cluster_cmd:
            cluster_param = f' --cluster "{cluster_cmd}"'
        else:
            cluster_log_dir = safe_mkdir(join(log_dir, 'cluster'))
            cluster_param = snakemake_utils.make_cluster_cmdl(cluster_log_dir, 'umccrise')

    ###############################
    #### Building command line ####
    ###############################

    cmd = ''
    if docker_wrapper_mode:
        cpus_avail = int(subprocess.check_output('docker system info | grep CPUs', shell=True).decode().split(': ')[1])
        if jobs > cpus_avail:
            click.echo(click.style(
                f'Warning: the number of jobs requested {jobs} is higher than then number of available CPUs '
                f'for the docker machine {cpus_avail}. Downgrading requested jobs to {cpus_avail}.', fg='red'))
            jobs = cpus_avail

        docker_tag = 'umccr/umccrise'
        if pcgr_data:
            docker_tag = 'umccr/umccrise_with_pcgr'

        cmd = (
            f'docker run -t ' +
            f'--cpus {jobs} ' +
            ''.join(f'-v={k}:{v} ' for k, v in docker_mounts.items()) +
            f'{docker_tag} '
        )

    snakefile = join(package_path(), 'Snakefile') if not docker_wrapper_mode else '/umccrise/umccrise/Snakefile'
    cmd += (
        f'snakemake '
        f'{" ".join(target_rule)} '
        f'--snakefile {snakefile} '
        f'--printshellcmds '
        f'--directory {output_dir if not docker_wrapper_mode else "/output_dir"} '
        f'-j {jobs} '
        f'--rerun-incomplete ' 
        f'--restart-times {restart_times} '
        f'{cluster_param} '
        f'--config {" ".join(k + "=" + v for k, v in conf.items())} '
    )

    #################
    #### Running ####
    #################

    if unlock:
        print('* Unlocking previous run... *')
        run_simple(cmd + ' --unlock')
        print('* Now rerunning *')

    try:
        run_simple(cmd)
    except subprocess.CalledProcessError:
        logger.error('--------')
        logger.error(f'Error running Umccrise: snakemake returned a non-zero status. Working directory: {output_dir}')
        run_simple(f'chmod -R a+r {log_dir}')
        if cluster_log_dir:
            logger.error(f'Review cluster job logs in {cluster_log_dir}')
        sys.exit(1)
    else:
        logger.error('--------')
        run_simple(f'chmod -R a+r {log_dir}')
        logger.info(f'Finished. Output directory: {output_dir}')


if __name__ == '__main__':
    main()
