import configparser
import functools
import itertools
import logging
import pathlib
import subprocess
from concurrent.futures import ProcessPoolExecutor, as_completed

# logger
log = logging.getLogger(__name__)
log.addHandler(logging.NullHandler())


def get_configuration(config_path):
    """
    Read .ini config file from given path
    """
    if isinstance(config_path, configparser.ConfigParser):
        return config_path
    ref_path_config = configparser.ConfigParser()
    ref_path_config.read(config_path)

    total_config = {}
    for name, section in ref_path_config.items():
        for k, v in section.items():
            total_config[k] = v
    return total_config


def test_cmd(tool_name, cmd_list):
    try:
        subprocess.run(cmd_list,
                       stdout=subprocess.PIPE,
                       stderr=subprocess.PIPE,
                       encoding='utf8',
                       check=True)
    except subprocess.CalledProcessError as e:
        log.error(f'Test {tool_name} got non-zero return code {e.returncode}')
        log.error(e.stderr)
        raise
    return


def valid_environments(config):
    log.info('Test mapping environments')

    # test cutadapt
    test_cmd(tool_name='cutadapt', cmd_list=['cutadapt', '--version'])
    # test samtools
    test_cmd(tool_name='samtools', cmd_list=['samtools', '--version'])
    # test picard, picard always have return code 1...
    test_cmd(tool_name='picard', cmd_list=['which', 'picard'])
    # test bismark_mapping
    test_cmd(tool_name='bismark_mapping', cmd_list=['bismark_mapping', '--version'])
    if config['mode'] != 'm3c':
        # test bowtie2
        test_cmd(tool_name='bowtie2', cmd_list=['bowtie2', '--version'])
    else:
        # test bowtie1
        test_cmd(tool_name='bowtie', cmd_list=['bowtie', '--version'])
    # test pigz
    test_cmd(tool_name='pigz', cmd_list=['pigz', '-V'])

    bismark_dir = pathlib.Path(config['bismark_reference'])
    if not bismark_dir.is_dir():
        raise TypeError(f"Bismark reference must be a directory contain a sub-dir named Bisulfite_Genome, "
                        f"generated by bismark_genome_preparation. Got a file path")
    if not bismark_dir.exists():
        raise FileNotFoundError(f"Bismark reference directory not found. "
                                f"Path in the config.ini is {bismark_dir}")

    allc_ref_fasta = pathlib.Path(config['reference_fasta'])
    allc_ref_fai = pathlib.Path(config['reference_fasta'] + '.fai')
    if not allc_ref_fasta.exists():
        raise FileNotFoundError(f"Reference fasta for ALLC generation not found. "
                                f"Path in the config.ini is {allc_ref_fasta}")
    if not allc_ref_fai.exists():
        raise FileNotFoundError(f".fai index for reference fasta not found. "
                                f"Path of fadix should be {allc_ref_fai}. "
                                f"You can use 'samtools fadix {allc_ref_fasta}' to generate.")
    return


def parse_index_fasta(fasta_path):
    records = {}
    with open(fasta_path) as f:
        key_line = True
        for line in f:
            if key_line:
                key = line.lstrip('>').rstrip('\n')
                key_line = False
            else:
                value = line.lstrip('^').rstrip('\n')
                records[key] = value
                key_line = True
    return records


def command_runner(commands, runner=None, cpu=1):
    if runner is None:
        from functools import partial
        runner = partial(subprocess.run,
                         stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE,
                         encoding='utf8',
                         shell=True,
                         check=True)
    if cpu <= 1:
        for command in commands:
            runner(command)
    else:
        with ProcessPoolExecutor(cpu) as pool:
            futures = []
            for command in commands:
                future = pool.submit(runner, command)
                futures.append(future)

            for future in as_completed(futures):
                try:
                    future.result()
                except subprocess.CalledProcessError as e:
                    print("Got error in fastq_qc, command was:")
                    print(command)
                    print(e.stdout)
                    print(e.stderr)
                    raise e
    return


def snakemake(workdir, snakefile, cores):
    try:
        subprocess.run([
            'snakemake', '-d', str(workdir), '--snakefile',
            str(snakefile), '--cores',
            str(cores)
        ],
            check=True,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            encoding='utf8')
    except subprocess.CalledProcessError as e:
        print(e.stdout)
        print(e.stderr)
        raise e
    return


def get_barcode_version(output_dir):
    fastq_dir = pathlib.Path(output_dir) / 'fastq'
    with open(fastq_dir / '.barcode_version') as f:
        return f.read()


def get_mode(output_dir):
    fastq_dir = pathlib.Path(output_dir) / 'fastq'
    with open(fastq_dir / '.mode') as f:
        return f.read()


MAPPING_MODE_CHOICES = ['mct', 'mc', 'm3c', '4m']
IUPAC_TABLE = {
    'A': 'A',
    'T': 'T',
    'C': 'C',
    'G': 'G',
    'R': 'AG',
    'Y': 'CT',
    'S': 'GC',
    'W': 'AT',
    'K': 'GT',
    'M': 'AC',
    'B': 'CGT',
    'D': 'AGT',
    'H': 'ATC',
    'V': 'ACG',
    'N': 'ATCGN'
}


@functools.lru_cache(maxsize=100)
def parse_mc_pattern(pattern: str) -> set:
    """
    parse mC context pattern
    """
    # IUPAC DNA abbr. table
    all_pos_list = []
    pattern = pattern.upper()
    for base in pattern:
        try:
            all_pos_list.append(IUPAC_TABLE[base])
        except KeyError:
            raise KeyError(f'Base {base} is not in IUPAC table.')
    context_set = set([''.join(i) for i in itertools.product(*all_pos_list)])
    return context_set
