# pylint: disable=invalid-name,protected-access
#!/usr/bin/env python3
# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.
"""
Utilities for downloading and building data.
These can be replaced if your particular file system does not support them.
"""
from typing import Union, List
from pathlib import Path
import time
import datetime
import os
import shutil
import requests

MODULE_ROOT = Path(__file__).parent.parent
FIXTURES_ROOT = (MODULE_ROOT / "tests" / "fixtures").resolve()
PACKAGE_ROOT = MODULE_ROOT.parent
DATA_ROOT = (PACKAGE_ROOT / "data").resolve()

class ProgressLogger(object):
    """Throttles and display progress in human readable form."""

    def __init__(self, throttle=1, should_humanize=True):
        """Initialize Progress logger.
        :param throttle: default 1, number in seconds to use as throttle rate
        :param should_humanize: default True, whether to humanize data units
        """
        self.latest = time.time()
        self.throttle_speed = throttle
        self.should_humanize = should_humanize

    def humanize(self, num, suffix='B'):
        """Convert units to more human-readable format."""
        if num < 0:
            return num
        for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
            if abs(num) < 1024.0:
                return "%3.1f%s%s" % (num, unit, suffix)
            num /= 1024.0
        return "%.1f%s%s" % (num, 'Yi', suffix)

    def log(self, curr, total, width=40, force=False):
        """Display a bar showing the current progress."""
        if curr == 0 and total == -1:
            print('[ no data received for this file ]', end='\r')
            return
        curr_time = time.time()
        if not force and curr_time - self.latest < self.throttle_speed:
            return
        else:
            self.latest = curr_time

        self.latest = curr_time
        done = min(curr * width // total, width)
        remain = width - done

        if self.should_humanize:
            curr = self.humanize(curr)
            total = self.humanize(total)

        progress = '[{}{}] {} / {}'.format(
            ''.join(['|'] * done),
            ''.join(['.'] * remain),
            curr,
            total
        )
        print(progress, end='\r')

def built(path, version_string=None):
    """Checks if '.built' flag has been set for that task.
    If a version_string is provided, this has to match, or the version
    is regarded as not built.
    """
    built_file_path = os.path.join(path, '.built')
    if not os.path.isfile(built_file_path):
        return False
    else:
        with open(built_file_path, 'r') as built_file:
            text = built_file.read().split('\n')
        if len(text) <= 2:
            return False
        for fname in text[1:-1]:
            if not os.path.isfile(os.path.join(path, fname)) and not os.path.isdir(os.path.join(path, fname)):
                return False
        return text[-1] == version_string if version_string else True


def mark_done(path, fnames, version_string='vXX'):
    """Marks the path as done by adding a '.built' file with the current
    timestamp plus a version description string if specified.
    """
    with open(os.path.join(path, '.built'), 'w') as built_file:
        built_file.write(str(datetime.datetime.today()))
        for fname in fnames:
            fname = fname.replace('.tar.gz', '').replace('.tgz', '').replace('.gz', '').replace('.zip', '')
            built_file.write('\n' + fname)
        built_file.write('\n' + version_string)


def download(url, path, fname, redownload=False):
    """Downloads file using `requests`. If ``redownload`` is set to false, then
    will not download tar file again if it is present (default ``True``)."""
    outfile = os.path.join(path, fname)
    curr_download = not os.path.isfile(outfile) or redownload
    print("[ downloading: " + url + " to " + outfile + " ]")
    retry = 5
    exp_backoff = [2 ** r for r in reversed(range(retry))]

    logger = ProgressLogger()

    while curr_download and retry >= 0:
        resume_file = outfile + '.part'
        resume = os.path.isfile(resume_file)
        if resume:
            resume_pos = os.path.getsize(resume_file)
            mode = 'ab'
        else:
            resume_pos = 0
            mode = 'wb'
        response = None

        with requests.Session() as session:
            try:
                header = {'Range': 'bytes=%d-' % resume_pos,
                          'Accept-Encoding': 'identity'} if resume else {}
                response = session.get(url, stream=True, timeout=5, headers=header)

                # negative reply could be 'none' or just missing
                if resume and response.headers.get('Accept-Ranges', 'none') == 'none':
                    resume_pos = 0
                    mode = 'wb'

                CHUNK_SIZE = 32768
                total_size = int(response.headers.get('Content-Length', -1))
                # server returns remaining size if resuming, so adjust total
                total_size += resume_pos
                done = resume_pos

                with open(resume_file, mode) as f:
                    for chunk in response.iter_content(CHUNK_SIZE):
                        if chunk:  # filter out keep-alive new chunks
                            f.write(chunk)
                        if total_size > 0:
                            done += len(chunk)
                            if total_size < done:
                                # don't freak out if content-length was too small
                                total_size = done
                            logger.log(done, total_size)
                    break
            except requests.exceptions.ConnectionError:
                retry -= 1
                # TODO Better way to clean progress bar?
                print(''.join([' '] * 60), end='\r')
                if retry >= 0:
                    print('Connection error, retrying. (%d retries left)' % retry)
                    time.sleep(exp_backoff[retry])
                else:
                    print('Retried too many times, stopped retrying.')
            finally:
                if response:
                    response.close()
    if retry < 0:
        raise RuntimeWarning('Connection broken too many times. Stopped retrying.')

    if curr_download and retry > 0:
        logger.log(done, total_size, force=True)
        print()
        if done < total_size:
            raise RuntimeWarning('Received less data than specified in ' +
                                 'Content-Length header for ' + url + '.' +
                                 ' There may be a download problem.')
        move(resume_file, outfile)


def make_dir(path):
    """Makes the directory and any nonexistent parent directories."""
    # the current working directory is a fine path
    if path != '':
        os.makedirs(path, exist_ok=True)


def move(path1, path2):
    """Renames the given file."""
    shutil.move(path1, path2)


def remove_dir(path):
    """Removes the given directory, if it exists."""
    shutil.rmtree(path, ignore_errors=True)


def untar(path, fname, deleteTar=True):
    """Unpacks the given archive file to the same directory, then (by default)
    deletes the archive file.
    """
    print('unpacking ' + fname)
    fullpath = os.path.join(path, fname)
    if '.tar.gz' in fname:
        shutil.unpack_archive(fullpath, path, format='gztar')
    else:
        shutil.unpack_archive(fullpath, path)
    if deleteTar:
        os.remove(fullpath)


def cat(file1, file2, outfile, deleteFiles=True):
    with open(outfile, 'wb') as wfd:
        for f in [file1, file2]:
            with open(f, 'rb') as fd:
                shutil.copyfileobj(fd, wfd, 1024 * 1024 * 10)
                # 10MB per writing chunk to avoid reading big file into memory.
    if deleteFiles:
        os.remove(file1)
        os.remove(file2)


def _get_confirm_token(response):
    for key, value in response.cookies.items():
        if key.startswith('download_warning'):
            return value
    return None

def download_from_google_drive(gd_id, destination):
    """Uses the requests package to download a file from Google Drive."""
    URL = 'https://docs.google.com/uc?export=download'

    with requests.Session() as session:
        response = session.get(URL, params={'id': gd_id}, stream=True)
        token = _get_confirm_token(response)

        if token:
            response.close()
            params = {'id': gd_id, 'confirm': token}
            response = session.get(URL, params=params, stream=True)

        CHUNK_SIZE = 32768
        with open(destination, 'wb') as f:
            for chunk in response.iter_content(CHUNK_SIZE):
                if chunk:  # filter out keep-alive new chunks
                    f.write(chunk)
        response.close()


def download_files(fnames: List[Union[str, Path]],
                   local_folder: str,
                   version: str = 'v1.0',
                   paths: Union[List[str], str] = 'aws') -> List[str]:
    r"""Download model/data files from a url.

    Args:
        fnames: List of filenames to download
        local_folder: Sub-folder of `./data` where models/data will
            be downloaded.
        version: Version of the model
        path: url or respective urls for downloading filenames.

    Return:
        List[str]: List of downloaded file path.
            If the downloaded file was a compressed file (`.tar.gz`,
            `.zip`, `.tgz`, `.gz`), return the path of the folder
            containing the extracted files.
    """

    dpath = str(DATA_ROOT / local_folder)
    out_paths = list(dpath + '/' + fname.replace('.tar.gz', '').replace('.tgz', '').replace('.gz', '').replace('.zip', '')
                     for fname in fnames)

    if not built(dpath, version):
        for fname in fnames:
            print('[building data: ' + dpath + '/' + fname + ']')
        if built(dpath):
            # An older version exists, so remove these outdated files.
            remove_dir(dpath)
        make_dir(dpath)

        if isinstance(paths, str):
            paths = [paths] * len(fnames)
        # Download the data.
        for fname, path in zip(fnames, paths):
            if path == 'aws':
                url = 'http://huggingface.co/downloads/models/'
                url += local_folder + '/'
                url += fname
            else:
                url = path + '/' + fname
            download(url, dpath, fname)
            if '.tar.gz' in fname or '.tgz' in fname or '.gz' in fname or '.zip' in fname:
                untar(dpath, fname)
        # Mark the data as built.
        mark_done(dpath, fnames, version)
    return out_paths
