# -*- coding: utf-8 -*-
# 版权所有 2020 深圳米筐科技有限公司（下称“米筐科技”）
#
# 除非遵守当前许可，否则不得使用本软件。
#
#     * 非商业用途（非商业用途指个人出于非商业目的使用本软件，或者高校、研究所等非营利机构出于教育、科研等目的使用本软件）：
#         遵守 Apache License 2.0（下称“Apache 2.0 许可”），您可以在以下位置获得 Apache 2.0 许可的副本：
#         http://www.apache.org/licenses/LICENSE-2.0。
#         除非法律有要求或以书面形式达成协议，否则本软件分发时需保持当前许可“原样”不变，且不得附加任何条件。
#
#     * 商业用途（商业用途指个人出于任何商业目的使用本软件，或者法人或其他组织出于任何目的使用本软件）：
#         未经米筐科技授权，任何个人不得出于任何商业目的使用本软件（包括但不限于向第三方提供、销售、出租、出借、转让本软件、本软件的衍生产品、引用或借鉴了本软件功能或源代码的产品或服务），任何法人或其他组织不得出于任何目的使用本软件，否则米筐科技有权追究相应的知识产权侵权责任。
#         在此前提下，对本软件的使用同样需要遵守 Apache 2.0 许可，Apache 2.0 许可与本许可冲突之处，以本许可为准。
#         详细的授权流程，请联系 public@ricequant.com 获取。
import os
import shutil
import tarfile
import tempfile
import time
import datetime
import dateutil
import sys

import click
import requests
import six
import h5py

from rqalpha.utils.i18n import gettext as _

from rqalpha.cmds.entry import cli
from rqalpha.utils import init_rqdatac_env


@cli.command(help=_("create bundle using RQDatac"))
@click.option('-d', '--data-bundle-path', default=os.path.expanduser('~/.rqalpha'), type=click.Path(file_okay=False))
@click.option("rqdatac_uri", '--rqdatac', '--rqdatac-uri', default=None,
              help='rqdatac uri, eg user:password or tcp://user:password@ip:port')
@click.option('--compression', default=False, is_flag=True, help='enable compression to reduce file size')
@click.option('-c', '--concurrency', type=click.INT, default=1)
def create_bundle(data_bundle_path, rqdatac_uri, compression, concurrency):
    try:
        import rqdatac
    except ImportError:
        click.echo(_(
            'rqdatac is required to create bundle. '
            'you can visit https://www.ricequant.com/welcome/rqdata to get rqdatac, '
            'or use "rqalpha download-bundle" to download monthly updated bundle.'
        ))
        return 1

    try:
        init_rqdatac_env(rqdatac_uri)
        rqdatac.init()
    except ValueError as e:
        click.echo(_('rqdatac init failed with error: {}').format(e))
        return 1

    os.makedirs(os.path.join(data_bundle_path, 'bundle'), exist_ok=True)
    from rqalpha.data.bundle import update_bundle as update_bundle_
    update_bundle_(os.path.join(data_bundle_path, 'bundle'), True, compression, concurrency)


@cli.command(help=_("Update bundle using RQDatac"))
@click.option('-d', '--data-bundle-path', default=os.path.expanduser('~/.rqalpha'), type=click.Path(file_okay=False))
@click.option("rqdatac_uri", '--rqdatac', '--rqdatac-uri', default=None,
              help='rqdatac uri, eg user:password or tcp://user:password@ip:port')
@click.option('--compression', default=False, type=click.BOOL, help='enable compression to reduce file size')
@click.option('-c', '--concurrency', type=click.INT, default=1)
def update_bundle(data_bundle_path, rqdatac_uri, compression, concurrency):
    try:
        import rqdatac
    except ImportError:
        click.echo(_(
            'rqdatac is required to update bundle. '
            'you can visit https://www.ricequant.com/welcome/rqdata to get rqdatac, '
            'or use "rqalpha download-bundle" to download monthly updated bundle.'
        ))
        return 1

    try:
        init_rqdatac_env(rqdatac_uri)
        rqdatac.init()
    except ValueError as e:
        click.echo(_('rqdatac init failed with error: {}').format(e))
        return 1

    if not os.path.exists(os.path.join(data_bundle_path, 'bundle')):
        click.echo(_('bundle not exist, use "rqalpha create-bundle" command instead'))
        return 1

    from rqalpha.data.bundle import update_bundle as update_bundle_
    succeed = update_bundle_(os.path.join(data_bundle_path, 'bundle'), False, compression, concurrency)
    if not succeed:
        sys.exit(1)


@cli.command(help=_("Download bundle (monthly updated)"))
@click.option('-d', '--data-bundle-path', default=os.path.expanduser('~/.rqalpha'), type=click.Path(file_okay=False))
@click.option('--confirm', default=True, is_flag=True)
def download_bundle(data_bundle_path, confirm):
    default_bundle_path = os.path.abspath(os.path.expanduser('~/.rqalpha/bundle'))
    if data_bundle_path is None:
        data_bundle_path = default_bundle_path
    else:
        data_bundle_path = os.path.abspath(os.path.join(data_bundle_path, './bundle/'))
    if (confirm and os.path.exists(data_bundle_path) and data_bundle_path != default_bundle_path and
            os.listdir(data_bundle_path)):
        click.confirm(_(u"""
    [WARNING]
    Target bundle path {data_bundle_path} is not empty.
    The content of this folder will be REMOVED before updating.
    Are you sure to continue?""").format(data_bundle_path=data_bundle_path), abort=True)

    tmp = os.path.join(tempfile.gettempdir(), 'rq.bundle')
    url, total_length = get_exactly_url()

    with open(tmp, 'wb') as out:
        download(out, total_length, url)

    shutil.rmtree(data_bundle_path, ignore_errors=True)
    os.makedirs(data_bundle_path)
    tar = tarfile.open(tmp, 'r:bz2')
    tar.extractall(data_bundle_path)
    tar.close()
    os.remove(tmp)
    six.print_(_(u"Data bundle download successfully in {bundle_path}").format(bundle_path=data_bundle_path))


@cli.command(help=_("Check bundle"))
@click.option('-d', '--data-bundle-path', default=os.path.expanduser('~/.rqalpha'), type=click.Path(file_okay=False))
def check_bundle(data_bundle_path):
    check_bundle_data(os.path.join(data_bundle_path, "bundle"))


CDN_URL = 'http://bundle.assets.ricequant.com/bundles_v4/rqbundle_%04d%02d.tar.bz2'


def get_exactly_url():
    day = datetime.date.today()
    proxy_uri = os.environ.get('RQALPHA_PROXY')
    while True:  # get exact url
        url = CDN_URL % (day.year, day.month)
        six.print_(_(u"try {} ...").format(url))
        r = requests.get(url, stream=True, proxies={'http': proxy_uri, 'https': proxy_uri})
        if r.status_code == 200:
            return url, int(r.headers.get('content-length'))

        day -= dateutil.relativedelta.relativedelta(months=1)


def download(out, total_length, url):
    retry_interval = 3
    retry_times = 5
    proxy_uri = os.environ.get('RQALPHA_PROXY')
    with click.progressbar(length=total_length, label=_(u"downloading ...")) as bar:
        for i in range(retry_times):
            try:
                headers = {'Range': "bytes={}-".format(bar.pos)}
                r = requests.get(url, headers=headers, stream=True, timeout=10, proxies={'http': proxy_uri,
                                                                                         'https': proxy_uri})
                for data in r.iter_content(chunk_size=8192):
                    bar.update(len(data))
                    out.write(data)

                if total_length == bar.pos:
                    return True  # Download complete . exit
            except requests.exceptions.RequestException:
                if i < retry_times - 1:
                    six.print_(_("\nDownload failed, retry in {} seconds.".format(retry_interval)))
                    time.sleep(retry_interval)
                else:
                    raise


def check_bundle_data(data_bundle_path):
    instruments = ["stocks", "indexes", "futures", "funds"]
    corrupt_files, not_exists_instruments = [], []
    for instrument in instruments:
        file = os.path.join(data_bundle_path, "{}.h5".format(instrument))
        if not os.path.exists(file):
            not_exists_instruments.append(instrument)
            continue
        try:
            with h5py.File(file, mode="r") as f:
                for order_book_id in f.keys():
                    __ = f[order_book_id][:1]
        except Exception:
            corrupt_files.append(file)
    if len(corrupt_files):
        click.echo("{}:\n{}".format(_("corrupted files"), corrupt_files))
        is_ok = input("{}(yes/no):".format(_("remove files"))).lower()
        if is_ok in ["yes", "y"]:
            [os.remove(file) for file in corrupt_files]
            click.echo(_("remove success"))
        elif is_ok in ["no", "n"]:
            click.echo(_("corrupted files not remove"))
        else:
            click.echo(_("input error"))
    elif len(not_exists_instruments):
        click.echo(_("bundle's day bar is incomplete, please update bundle"))
    else:
        click.echo(_("good bundle's day bar"))

