import re
import subprocess
import time
import os
import shutil
from multiprocessing import Pool, TimeoutError
from multiprocessing.dummy import Pool as ThreadPool
from functools import partial
from lib.common.util import printer_list, logging_init, time_now

__author__ = 'Yuxing Xu'


def mkdir(dir_name, keep=False):
    if keep is False:
        if os.path.exists(dir_name):
            shutil.rmtree(dir_name)
        os.makedirs(dir_name)
    else:
        if not os.path.exists(dir_name):
            os.makedirs(dir_name)

    return dir_name


def rmdir(dir_name):
    if os.path.exists(dir_name):
        if os.path.isdir(dir_name):
            shutil.rmtree(dir_name)
        else:
            os.remove(dir_name)


def have_file(file_name, null_is_ok=False):
    if os.path.exists(file_name):
        if null_is_ok:
            return True
        elif os.path.getsize(file_name):
            return True
    else:
        return False


def get_file_name(give_path):
    return give_path.split("/")[-1]


def get_file_dir(file_path):
    abs_file_path = os.path.abspath(file_path)
    return os.path.dirname(abs_file_path)


def copy_file(source_path, target_path, keep=False):
    source_path = os.path.abspath(source_path)
    target_path = os.path.abspath(target_path)

    new_path = shutil.copy(source_path, target_path)

    return new_path


def ln_file(source_path, target_path):

    if os.path.isdir(target_path):
        target_path = target_path + "/" + get_file_name(source_path)

    os.symlink(os.path.abspath(source_path), target_path)

    return target_path


def move_file(source_path, target_path, keep=False):
    source_path = os.path.abspath(source_path)
    target_path = os.path.abspath(target_path)

    new_path = shutil.move(source_path, target_path)

    return new_path


def gunzip_file(raw_file):
    raw_file = os.path.abspath(raw_file)

    cmd_string = "gunzip " + raw_file
    cmd_run(cmd_string, silence=True)

    gunzip_file = re.sub(".gz$", "", raw_file)

    return gunzip_file


def gzip_file(raw_file):
    raw_file = os.path.abspath(raw_file)

    cmd_string = "gzip " + raw_file
    cmd_run(cmd_string, silence=True)

    gzip_file = raw_file + ".gz"

    return gzip_file


def md5sum_check(file_name, original_md5):
    import hashlib

    if not os.path.exists(file_name):
        return False

    with open(file_name, "rb") as file_to_check:
        data = file_to_check.read()
        md5_returned = hashlib.md5(data).hexdigest()

    if original_md5 == md5_returned:
        return True
    else:
        return False


def remove_file_name_suffix(file_name, subsuffix_level=0):
    """
    remove the suffix for a file name
    :param file_name: give me a file name, like: "/home/xuyuxing/file.txt" or "~/work/file.txt" or just "file.txt"
    :param subsuffix_level: sometimes a file can have more than one suffix, like "~/work/file.txt.gz",
           how many you want to remove, 0 meanings remove all suffix
    :return: file without suffix, will keep dir_name as a path, to get base file name: os.path.basename(file_name)
    """

    file_name = os.path.abspath(file_name)
    dir_name = os.path.dirname(file_name)
    base_name = os.path.basename(file_name)
    base_name_split = base_name.split(".")
    if subsuffix_level == 0:
        subsuffix_level = len(base_name_split) - 1
    elif subsuffix_level > len(base_name_split) - 1:
        raise ValueError("There are not enough suffix level to remove")

    return dir_name + "/" + printer_list(base_name.split(".")[:-subsuffix_level], ".")


def md5sum_maker(file_name):
    import hashlib

    if not os.path.exists(file_name):
        FileNotFoundError("No such file or directory: %s" % file_name)

    with open(file_name, "rb") as file_to_check:
        data = file_to_check.read()
        md5_returned = hashlib.md5(data).hexdigest()

    return md5_returned


def cmd_run(cmd_string, cwd=None, retry_max=5, silence=False, log_file=None):
    module_logger = logging_init("cmd_run", log_file)
    module_logger.info("Calling a bash cmd with retry_max %d: %s" %
                       (retry_max, cmd_string))
    if not silence:
        print("Running " + str(retry_max) + " " + cmd_string)
    p = subprocess.Popen(cmd_string, shell=True,
                         stderr=subprocess.PIPE, stdout=subprocess.PIPE, cwd=cwd)
    output, error = p.communicate()
    if not silence:
        print(error.decode())
    returncode = p.poll()
    module_logger.info("Finished bash cmd with returncode as %d" % returncode)
    if returncode == 1:
        if retry_max > 1:
            retry_max = retry_max - 1
            cmd_run(cmd_string, cwd=cwd, retry_max=retry_max)
    del module_logger.handlers[:]

    output = output.decode()
    error = error.decode()

    return (not returncode, output, error)


def multiprocess_running(f, args_list, pool_num, log_file=None, silence=True, args_id_list=None, timeout=None):
    """
    :param f: function name
    :param args_list: a list include every args in a tuple
    :param pool_num: process num of running at sametime
    :param silence:
    :return:

    example

    def f(x, y):
        time.sleep(1)
        return x * y

    args_list = list(zip(range(1,200),range(2,201)))
    pool_num = 5

    multiprocess_running(f, args_list, pool_num, log_file='/lustre/home/xuyuxing/Work/Other/saif/Meth/tmp/log')

    """
    num_tasks = len(list(args_list))

    module_log = logging_init(f.__name__, log_file)
    if not silence:
        print('args_list have %d object and pool_num is %d' %
              (num_tasks, pool_num))
    module_log.info('running with mulitprocess')
    module_log.info('args_list have %d object and pool_num is %d' %
                    (num_tasks, pool_num))

    p_dict = {}

    if len(args_list) > 0:

        if timeout is None:

            f_with_para = partial(get_more_para, f)

            #    def get_more_para(para_tuple):
            #        return f(*para_tuple)

            #    args_list_unziped = zip(*args_list)
            start_time = time.time()
            if not silence:
                print(time_now() + '\tBegin: ')
            module_log.info('Begin: ')
            with Pool(processes=pool_num) as pool:
                # for i, output in enumerate(pool.imap_unordered(f_with_para, args_list, chunksize=1)):
                for i, output in enumerate(pool.imap(f_with_para, args_list, chunksize=1)):
                    if args_id_list:
                        job_id = args_id_list[i]
                    else:
                        job_id = 'ID_' + str(i)

                    p_dict[job_id] = {
                        'args': args_list[i],
                        'output': output,
                        'error': None
                    }
                    # print(i)
                    round_time = time.time()
                    if round_time - start_time > 5:
                        if not silence:
                            print(time_now() + '\t%d/%d %.2f%% parsed' %
                                  (i, num_tasks, i / num_tasks * 100))
                        module_log.info('%d/%d %.2f%% parsed' %
                                        (i, num_tasks, i / num_tasks * 100))
                        start_time = round_time
                    # module_log.info('%d/%d %.2f%% parsed' % (i, num_tasks, i / num_tasks * 100))

            module_log.info('%d/%d %.2f%% parsed' %
                            (i, num_tasks, i / num_tasks * 100))
            module_log.info('All args_list task finished')

            if not silence:
                print(time_now() + '\t%d/%d %.2f%% parsed' %
                      (i, num_tasks, i / num_tasks * 100))
                print(time_now() + '\tAll args_list task finished')

        else:

            f_with_para = partial(get_more_para, f)
            abortable_func = partial(
                abortable_worker, f_with_para, timeout=timeout)

            #    def get_more_para(para_tuple):
            #        return f(*para_tuple)

            #    args_list_unziped = zip(*args_list)

            start_time = time.time()

            if not silence:
                print(time_now() + '\tBegin: ')
            module_log.info('Begin: ')

            with Pool(processes=pool_num, maxtasksperchild=1) as pool:
                # for i, output in enumerate(pool.imap_unordered(f_with_para, args_list, chunksize=1)):
                it = pool.imap(abortable_func, args_list, chunksize=1)
                i = -1
                while 1:
                    i += 1
                    # print(i)

                    try:
                        output = it.next()

                        if args_id_list:
                            job_id = args_id_list[i]
                        else:
                            job_id = 'ID_' + str(i)

                        if output == "Aborting due to timeout":
                            p_dict[job_id] = {
                                'args': args_list[i],
                                'output': None,
                                'error': "timeout"
                            }
                        else:
                            p_dict[job_id] = {
                                'args': args_list[i],
                                'output': output,
                                'error': None
                            }

                        # print(i)
                        round_time = time.time()
                        if round_time - start_time > 5:
                            if not silence:
                                print('%d/%d %.2f%% parsed' %
                                      (i, num_tasks, i / num_tasks * 100))
                            module_log.info('%d/%d %.2f%% parsed' %
                                            (i, num_tasks, i / num_tasks * 100))
                            start_time = round_time

                        # module_log.info('%d/%d %.2f%% parsed' % (i, num_tasks, i / num_tasks * 100))

                    except StopIteration:
                        break

            module_log.info('%d/%d %.2f%% parsed' %
                            (i, num_tasks, i / num_tasks * 100))
            module_log.info('All args_list task finished')

            if not silence:
                print(time_now() + '%d/%d %.2f%% parsed' %
                      (i, num_tasks, i / num_tasks * 100))
                print(time_now() + 'All args_list task finished')

    del module_log.handlers[:]

    return p_dict


def get_more_para(f, para_tuple):
    return f(*para_tuple)


def abortable_worker(func, *args, **kwargs):
    timeout = kwargs.get('timeout', None)
    p = ThreadPool(1)
    res = p.apply_async(func, args=args)
    try:
        out = res.get(timeout)  # Wait timeout seconds for func to complete.
        return out
    except TimeoutError:
        return "Aborting due to timeout"
        raise
