import sys
import io
import os
import os.path as op
import logging
import json
import traceback
import shutil
import time
from urllib.request import urlopen, Request
from urllib.request import HTTPError
import numpy as np
import subprocess as sp
from pathos.multiprocessing import ProcessingPool as Pool


def setup_logger(log_file: str = '', logging_level=logging.INFO):
    logging.getLogger().setLevel(logging_level)
    log_format = "%(asctime)s - [%(levelname)s] - [P:%(process)d] - %(message)s"
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(logging.Formatter(fmt=log_format))
    logging.getLogger().addHandler(console_handler)

    if log_file != '':
        file_handler = logging.FileHandler(log_file, mode='w')
        file_handler.setFormatter(logging.Formatter(fmt=log_format))
        logging.getLogger().addHandler(file_handler)


def get_logger(name: str, log_file: str = '', logging_level=logging.INFO):
    logger = logging.getLogger(name)
    logger.setLevel(logging_level)
    log_format = "%(asctime)s - [%(levelname)s] - [P:%(process)d] - %(message)s"

    if log_file != '':
        file_handler = logging.FileHandler(log_file, mode='w')
        file_handler.setFormatter(logging.Formatter(fmt=log_format))
        logger.addHandler(file_handler)
    else:
        console_handler = logging.StreamHandler(sys.stdout)
        console_handler.setFormatter(logging.Formatter(fmt=log_format))
        logger.addHandler(console_handler)

    return logger


def read_lines(filename, encoding=None):
    """Helper function to read lines from a text file.

    Args:
        filename (str): filename to be read.
        encoding (str): character encoding, e.g. 'utf-8'
    """
    with io.open(filename, 'r', encoding=encoding) as fp:
        for line in fp:
            yield line.strip()


def file_exists(path, exit_if_not_exists=True):
    """Helper function to check existence of the given file path.

    Args:
        path (str): file path to check.

    Returns:
        bool: True if file exists, else False
    """
    if path is not None and len(path) > 0:
        assert not op.isdir(path), '{} is a folder'.format(path)
        if not os.path.isfile(path):
            if exit_if_not_exists:
                logging.error("%s not exists!", path)
                sys.exit(1)
            else:
                return False
        else:
            return True


def folder_exists(path, exit_if_not_exists=True):
    """Helper function to check existence of the given folder path.

    Args:
        path (str): folder path to check.

    Returns:
        bool: True if path exists, else False
    """
    if path == '' or path == '.':
        return True
    if path is not None and len(path) > 0:
        assert not op.isfile(path), '{} is a file'.format(path)
        if not os.path.exists(path) and not op.islink(path):
            if exit_if_not_exists:
                logging.error("%s is not a folder!", path)
                sys.exit(1)
            else:
                return False
        else:
            return True


def ensure_directory(path):
    """Check exsitence of the given directory path. If not, create a new directory.

    Args:
        path (str): path of a given directory.
    """
    if path == '' or path == '.':
        return
    if path is not None and len(path) > 0:
        assert not op.isfile(path), '{} is a file'.format(path)
        if not os.path.exists(path) and not op.islink(path):
            try:
                os.makedirs(path)
            except:
                if os.path.isdir(path):
                    # another process has done makedir
                    pass
                else:
                    raise
        # we should always check if it succeeds.
        assert op.isdir(op.abspath(path)), path


def image_url_to_bytes(url):
    req = Request(url, headers={
        "User-Agent": "Mozilla/5.0 (X11; Linux i686) AppleWebKit/537.17 (KHTML, like Gecko) Chrome/24.0.1312.27 Safari/537.17"})
    try:
        response = urlopen(req, None, 10)
        if response.code != 200:
            logging.info("url: {}, error code: {}".format(url, response.code))
            return None
        data = response.read()
        response.close()
        return data
    except Exception as e:
        logging.info("error downloading: {}".format(e))
    return None


def concat_files(ins, out):
    ensure_directory(op.dirname(out))
    out_tmp = out + '.tmp'
    with open(out_tmp, 'wb') as fp_out:
        for i, f in enumerate(ins, 1):
            print('concating {}/{} - {}'.format(i, len(ins), f))
            with open(f, 'rb') as fp_in:
                shutil.copyfileobj(fp_in, fp_out, 1024 * 1024 * 10)
    os.rename(out_tmp, out)


def append_files(src, append_list):
    with open(src, 'ab') as fp_out:
        for i, f in enumerate(append_list, 1):
            print('appending {}/{} - {}'.format(i, len(append_list), f))
            with open(f, 'rb') as fp_in:
                shutil.copyfileobj(fp_in, fp_out, 1024 * 1024 * 10)


def url_to_bytes(url):
    try:
        fp = urlopen(url, timeout=30)
        buf = fp.read()
        real_url = fp.geturl()
        if real_url != url and (not real_url.startswith('https') or
                                real_url.replace('https', 'http') != url):
            logging.info('new url = {}; old = {}'.format(fp.geturl(), url))
            # the image gets redirected, which means the image is not available
            return None
        return buf
    except HTTPError as err:
        logging.error("url: {}; error code {}; message: {}".format(
            url, err.code, err.msg))
        return None
    except:
        logging.error("url: {}; unknown {}".format(
            url, traceback.format_exc()))
        return None


def json_dump(obj):
    # order the keys so that each operation is deterministic though it might be
    # slower
    return json.dumps(obj, sort_keys=True, separators=(',', ':'))


def parallel_map(func, all_task, num_worker=16):
    if num_worker > 0:
        m = Pool(num_worker)
        return m.map(func, all_task)
    else:
        result = []
        for t in all_task:
            result.append(func(t))
        return result


def print_trace():
    traceback.print_exc()


def try_once(func):
    def func_wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            logging.info('ignore error \n{}'.format(str(e)))
            print_trace()

    return func_wrapper


@try_once
def try_delete(f):
    os.remove(f)


def write_to_file(contxt, file_name, append=False):
    p = os.path.dirname(file_name)
    ensure_directory(p)
    if type(contxt) is str:
        contxt = contxt.encode()
    flag = 'wb'
    if append:
        flag = 'ab'
    with open(file_name, flag) as fp:
        fp.write(contxt)


def copy_file(src, dest):
    tmp = dest + '.tmp'
    # we use rsync because it could output the progress
    cmd_run('rsync {} {} --progress'.format(src, tmp).split(' '))
    os.rename(tmp, dest)


def ensure_copy_file(src, dest):
    ensure_directory(op.dirname(dest))
    if not op.isfile(dest):
        copy_file(src, dest)


def limited_retry_agent(num, func, *args, **kwargs):
    for i in range(num):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            logging.info('fails with \n{}: tried {}-th time'.format(
                e,
                i + 1))
            print_trace()
            if i == num - 1:
                raise
            time.sleep(5)


def cmd_run(list_cmd, return_output=False, env=None,
            working_dir=None,
            stdin=sp.PIPE,
            shell=False,
            dry_run=False,
            ):
    logging.info('start to cmd run: {}'.format(' '.join(map(str, list_cmd))))
    # if we dont' set stdin as sp.PIPE, it will complain the stdin is not a tty
    # device. Maybe, the reson is it is inside another process.
    # if stdout=sp.PIPE, it will not print the result in the screen
    e = os.environ.copy()
    if 'SSH_AUTH_SOCK' in e:
        del e['SSH_AUTH_SOCK']
    if working_dir:
        ensure_directory(working_dir)
    if env:
        for k in env:
            e[k] = env[k]
    if dry_run:
        # we need the log result. Thus, we do not return at teh very beginning
        return
    if not return_output:
        # if env is None:
        # p = sp.Popen(list_cmd, stdin=sp.PIPE, cwd=working_dir)
        # else:
        if shell:
            p = sp.Popen(' '.join(list_cmd),
                         stdin=stdin,
                         env=e,
                         cwd=working_dir,
                         shell=True)
        else:
            p = sp.Popen(list_cmd,
                         stdin=sp.PIPE,
                         env=e,
                         cwd=working_dir)
        message = p.communicate()
        if p.returncode != 0:
            raise ValueError(message)
    else:
        if shell:
            message = sp.check_output(' '.join(list_cmd),
                                      env=e,
                                      cwd=working_dir,
                                      shell=True)
        else:
            message = sp.check_output(list_cmd,
                                      env=e,
                                      cwd=working_dir)
        logging.info('finished the cmd run')
        return message.decode('utf-8')


def get_file_size(f):
    if not op.isfile(f):
        return 0
    return os.stat(f).st_size
