# -*- coding: utf-8 -*-
# Copyright (C) 2020 Unbabel
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import os
import subprocess
import urllib.request
import zipfile
from pathlib import Path
from typing import List
from urllib.parse import urlparse

from tqdm import tqdm

logger = logging.getLogger(__name__)


available_legacy_metrics = {
    # WMT20 Models
    "emnlp20-comet-rank": "https://unbabel-experimental-models.s3.amazonaws.com/comet/wmt20/emnlp20-comet-rank.tar.gz",
    "wmt20-comet-da": "https://unbabel-experimental-models.s3.amazonaws.com/comet/wmt20/wmt20-comet-da.tar.gz",
    "wmt20-comet-qe-da": "https://unbabel-experimental-models.s3.amazonaws.com/comet/wmt20/wmt20-comet-qe-da.tar.gz",
    "wmt20-comet-qe-da-v2": "https://unbabel-experimental-models.s3.amazonaws.com/comet/wmt20/wmt20-comet-qe-da-v2.tar.gz",

    # WMT21 Models
    "wmt21-comet-da": "https://unbabel-experimental-models.s3.amazonaws.com/comet/wmt21/wmt21-comet-da.tar.gz",
    "wmt21-comet-mqm": "https://unbabel-experimental-models.s3.amazonaws.com/comet/wmt21/wmt21-comet-mqm.tar.gz",
    "wmt21-cometinho-mqm": "https://unbabel-experimental-models.s3.amazonaws.com/comet/wmt21/wmt21-cometinho-mqm.tar.gz",
    "wmt21-cometinho-da": "https://unbabel-experimental-models.s3.amazonaws.com/comet/wmt21/wmt21-cometinho-da.tar.gz",
    "wmt21-comet-qe-mqm": "https://unbabel-experimental-models.s3.amazonaws.com/comet/wmt21/wmt21-comet-qe-mqm.tar.gz",
    "wmt21-comet-qe-da": "https://unbabel-experimental-models.s3.amazonaws.com/comet/wmt21/wmt21-comet-qe-da.tar.gz",

    # EAMT22 Models
    "eamt22-cometinho-da": "https://unbabel-experimental-models.s3.amazonaws.com/comet/eamt22/eamt22-cometinho-da.tar.gz",
    "eamt22-prune-comet-da": "https://unbabel-experimental-models.s3.amazonaws.com/comet/eamt22/eamt22-prune-comet-da.tar.gz",
}


def get_cache_folder():
    cache_directory = Path.home() / ".cache" / "torch" / "unbabel_comet"
    if not cache_directory.exists():
        cache_directory.mkdir(exist_ok=True, parents=True)

    return str(cache_directory)


def _reporthook(t):
    """``reporthook`` to use with ``urllib.request`` that prints the
        process of the download.

    Uses ``tqdm`` for progress bar.

    **Reference:**
    https://github.com/tqdm/tqdm

    """
    last_b = [0]

    def inner(b: int = 1, bsize: int = 1, tsize: int = None):
        """
        :param b: Number of blocks just transferred [default: 1].
        :param bsize: Size of each block (in tqdm units) [default: 1].
        :param tsize: Total size (in tqdm units).
            If [default: None] remains unchanged.
        """
        if tsize is not None:
            t.total = tsize
        t.update((b - last_b[0]) * bsize)
        last_b[0] = b

    return inner


def _maybe_extract(compressed_filename: str, directory: str, extension: str = None):
    """Extract a compressed file to ``directory``.

    :param compressed_filename: Compressed file.
    :param directory: Extract to directory.
    :param extension: Extension of the file; Otherwise, attempts to
        extract extension from the filename.
    """
    logger.info("Extracting {}".format(compressed_filename))

    if extension is None:
        basename = os.path.basename(compressed_filename)
        extension = basename.split(".", 1)[1]

    if "zip" in extension:
        with zipfile.ZipFile(compressed_filename, "r") as zip_:
            zip_.extractall(directory)

    elif "tar.gz" in extension or "tgz" in extension:
        # `tar` is much faster than python's `tarfile` implementation
        with open(os.devnull, "w") as devnull:
            subprocess.call(
                ["tar", "-C", directory, "-zxvf", compressed_filename], stdout=devnull
            )

    elif "tar" in extension:
        with open(os.devnull, "w") as devnull:
            subprocess.call(
                ["tar", "-C", directory, "-xvf", compressed_filename], stdout=devnull
            )

    logger.info("Extracted {}".format(compressed_filename))


def _get_filename_from_url(url):
    """Return a filename from a URL

    Args:
        url (str): URL to extract filename from

    Returns:
        (str): Filename in URL
    """
    parse = urlparse(url)
    return os.path.basename(parse.path)


def _check_download(*filepaths):
    """Check if the downloaded files are found.

    Args:
        filepaths (list of str): Check if these filepaths exist

    Returns:
        (bool): Returns True if all filepaths exist
    """
    return all([os.path.isfile(filepath) for filepath in filepaths])


def download_file_maybe_extract(
    url: str,
    directory: str,
    filename: str = None,
    extension: str = None,
    check_files: List[str] = None,
):
    """Download the file at ``url`` to ``directory``.
        Extract to ``directory`` if tar or zip.

    :param url: Url of file (str or Path).
    :param directory: Directory to download to.
    :param filename: Name of the file to download; Otherwise, a filename is extracted
        from the url.
    :param extension: Extension of the file; Otherwise, attempts to extract extension
        from the filename.
    :param check_files: Check if these files exist, ensuring the download
        succeeded. If these files exist before the download, the download is skipped.

    :return: Filename of download file.
    """
    if filename is None:
        filename = _get_filename_from_url(url)

    check_files = [] if not check_files else check_files

    directory = str(directory)
    filepath = os.path.join(directory, filename)
    check_files = [os.path.join(directory, str(f)) for f in check_files]

    if len(check_files) > 0 and _check_download(*check_files):
        return filepath

    if not os.path.isdir(directory):
        os.makedirs(directory)

    logger.info("Downloading {}".format(filename))

    # Download
    with tqdm(unit="B", unit_scale=True, miniters=1, desc=filename) as t:
        urllib.request.urlretrieve(url, filename=filepath, reporthook=_reporthook(t))

    _maybe_extract(
        compressed_filename=filepath, directory=directory, extension=extension
    )

    if not _check_download(*check_files):
        raise ValueError("[DOWNLOAD FAILED] `*check_files` not found")

    return filepath


def download_model_legacy(model: str, saving_directory: str = None) -> str:
    """
    Function that loads legacy pretrained models from AWS that are not on
     Hugging Face hub.

    :param model: Name of the model to be loaded.
    :param saving_directory: RELATIVE path to the saving folder (must end with /).

    Return:
        - Path to model checkpoint.
    """

    if saving_directory is None:
        saving_directory = get_cache_folder()

    if not saving_directory.endswith("/"):
        saving_directory += "/"

    if not os.path.exists(saving_directory):
        os.makedirs(saving_directory)

    if os.path.isdir(saving_directory + model):
        logger.info(f"{model} is already in cache.")
        if not model.endswith("/"):
            model += "/"

    elif model not in available_legacy_metrics.keys():
        raise Exception(
            f"{model} is not in the `available_legacy_metrics` or is a valid checkpoint folder."
        )

    elif available_legacy_metrics[model].startswith("https://"):
        download_file_maybe_extract(
            available_legacy_metrics[model], directory=saving_directory
        )

    else:
        raise Exception("Invalid model name!")

    # CLEAN Cache
    if os.path.exists(saving_directory + model + ".zip"):
        os.remove(saving_directory + model + ".zip")
    if os.path.exists(saving_directory + model + ".tar.gz"):
        os.remove(saving_directory + model + ".tar.gz")
    if os.path.exists(saving_directory + model + ".tar"):
        os.remove(saving_directory + model + ".tar")

    checkpoints_folder = saving_directory + model + "/checkpoints"
    checkpoints = [
        file for file in os.listdir(checkpoints_folder) if file.endswith(".ckpt")
    ]
    checkpoint = checkpoints[-1]
    checkpoint_path = checkpoints_folder + "/" + checkpoint
    return checkpoint_path
