#  ******************************************************************************
#  Copyright (c) 2020 Konduit K.K.
#
#  This program and the accompanying materials are made available under the
#  terms of the Apache License, Version 2.0 which is available at
#  https://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.
#
#  SPDX-License-Identifier: Apache-2.0
#  *****************************************************************************

import os
import sys
import subprocess
import requests
from hurry.filesize import size
from packaging.version import parse
import click

USER_PATH = os.path.expanduser("~")
KONDUIT_BASE_DIR = os.path.join(USER_PATH, ".konduit-serving")
KONDUIT_SOURCE_DIR = os.path.join(KONDUIT_BASE_DIR, "source")
KONDUIT_JAR_DIR = os.path.join(KONDUIT_BASE_DIR, "jar")
KONDUIT_JAR_PATH = os.path.join(KONDUIT_JAR_DIR, "konduit.jar")

INCOMPATIBLE_COMPILATION_TAGS = ["cli_base", "cli_base_2", "cli_base_3", "cli_base_4", "cli_base_5"]
DOWNLOAD_TAG = "cli_base"

LAST_COMPATIBLE_KONDUIT_VERSION = "0.1.0-SNAPSHOT"
DEFAULT_KONDUIT_TAG = "cli_base_5"
KONDUIT_JAR_URL_FORMAT = "https://github.com/KonduitAI/konduit-serving/releases/download/" \
                         "{tag}/konduit-serving-uberjar-{version}-{spin}-{platform}-{chip}${cuda_version}.jar"

os_choices = [
    "windows-x86_64",
    "linux-x86_64",
    "macosx-x86_64",
    "linux-armhf",
]


def get_platform():
    if sys.platform.startswith("win32"):
        return "windows-x86_64"
    elif sys.platform.startswith("darwin"):
        return "macosx-x86_64"
    elif sys.platform.startswith("linux"):
        return "linux-x86_64"
    else:
        raise RuntimeError("Please specify '--os'. Possible values are: " + os_choices.__str__())


def download_if_required(url, save_path):
    if not os.path.exists(KONDUIT_JAR_DIR):
        os.makedirs(KONDUIT_JAR_DIR)

    response = requests.get(url, stream=True)
    status_code = response.status_code
    if status_code != 200:
        print("Download path: {}".format(url))
        print("Failed with a status code of {}".format(status_code))
        return

    total = response.headers.get('content-length')

    if os.path.exists(save_path) and os.path.getsize(KONDUIT_JAR_PATH) == total:
        print("The required CLI binary has already been downloaded.")
        return
    else:
        print("Downloading command line binaries from " + url)

    with open(save_path, 'wb') as f:
        if total is None:
            f.write(response.content)
        else:
            downloaded = 0
            total = int(total)
            for data in response.iter_content(chunk_size=max(int(total / 1000), 1024 * 1024)):
                downloaded += len(data)
                f.write(data)
                done = int(50 * downloaded / total)
                sys.stdout.write('\r[{}{}]'.format('█' * done, '.' * (50 - done)) +
                                 (" ({}/{})".format(size(downloaded), size(total))))
                sys.stdout.flush()
    sys.stdout.write('\n')


def git_clone_konduit(use_https=True, tag=DEFAULT_KONDUIT_TAG):
    """Clone the konduit-serving git repo, if it doesn't already exist locally."""

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

    if use_https:
        repo = "https://github.com/KonduitAI/konduit-serving.git"
    else:
        repo = "git@github.com:KonduitAI/konduit-serving.git"
    try:
        if not os.listdir(KONDUIT_SOURCE_DIR):
            subprocess.call(["git", "clone", repo, KONDUIT_SOURCE_DIR],
                            shell=sys.platform.startswith("win"))
        # Pulling in changes if needed
        subprocess.call(["git", "-C", KONDUIT_SOURCE_DIR, "fetch", "--all"])
        subprocess.call(["git", "checkout", tag], cwd=KONDUIT_SOURCE_DIR,
                        shell=sys.platform.startswith("win"))
        subprocess.call(["git", "-C", KONDUIT_SOURCE_DIR, "pull"])
    except Exception as e:
        raise RuntimeError(">>> Could not clone konduit-serving repository and switch to commit hash {}"
                           "Make sure to have git installed. Type 'konduit-init --help' for help .\n"
                           .format(tag), e)


def build_jar(operating_sys, spin, chip, cuda_version, options):
    """Build the actual JAR, using our mvnw wrapper under the hood."""

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

    if operating_sys is None:
        operating_sys = get_platform()

    # Building the uber-jar file
    try:
        command = [
            "python",
            os.path.join(KONDUIT_SOURCE_DIR, "build_jar.py"),
            "--os",
            operating_sys,
            "--source",
            KONDUIT_SOURCE_DIR,
            "--spin",
            spin,
            "--chip",
            chip,
            "--cuda-version",
            cuda_version,
            "--target",
            KONDUIT_JAR_PATH
        ]
        command.extend(["--option {}".format(option) for option in options])

        subprocess.call(
            command,
            shell=sys.platform.startswith("win")
        )
    except Exception as e:
        RuntimeError("Failed to build jar.\n", e)


def get_git_tags():
    response = requests.get("https://api.github.com/repos/KonduitAI/konduit-serving/tags")
    status_code = response.status_code

    if status_code == 200:
        tags = response.json()
        return [tag['name'] for tag in tags]
    else:
        return [DEFAULT_KONDUIT_TAG]


def get_jar_url(platform, version, spin, chip, cuda_version=None):
    if not platform:
        platform = get_platform()

    return KONDUIT_JAR_URL_FORMAT.format(version=version,
                                         platform=platform,
                                         tag=DOWNLOAD_TAG,
                                         spin=spin,
                                         chip=chip,
                                         cuda_version=("-cuda" + cuda_version) if chip == "gpu" else "")


git_tags = list(set(get_git_tags()).difference(INCOMPATIBLE_COMPILATION_TAGS))
if len(git_tags) == 0:
    git_tags = ["master"]
DEFAULT_KONDUIT_TAG = git_tags[0]  # Assuming the first one in the response is the most recent one


@click.command()
@click.option(
    "-p", "--platform",
    help="Your operating system. "
         "Defaults to the cpu version of the "
         "current OS platform in use.",
    type=click.Choice(os_choices),
)
@click.option(
    "--https",
    type=bool,
    default=True,
    show_default=True,
    help="If True, use HTTPS to clone konduit-serving, else SSH."
)
@click.option(
    "-t", "--tag",
    default=DEFAULT_KONDUIT_TAG,
    show_default=True,
    help="The git tag to build the source code from.",
    type=click.Choice(git_tags)
)
@click.option(
    "-s", "--spin",
    default="all",
    show_default=True,
    help="Whether to bundle Python ('python'), PMML ('pmml'), both ('all') " +
         "or neither ('minimal'). Python bundling is not encouraged with ARM, " +
         "and PMML bundling is not encouraged if agpl license is an issue."
)
@click.option(
    "-c", "--chip",
    default="cpu",
    show_default=True,
    help="Specifies the chip architecture which could be cpu, gpu (CUDA) or arm.",
    type=click.Choice(["cpu", "gpu", "arm"])
)
@click.option(
    "-cv", "--cuda-version",
    default="10.2",
    show_default=True,
    help="Specifies the CUDA version to compile when `--chip == gpu`",
    type=click.Choice(["10.0", "10.1", "10.2"])
)
@click.option(
    "-o", "--options",
    multiple=True,
    help="Extra java property to add to the build command. Can be used multiple times."
         "Format: -o <property1.name>=<property1.value>"
)
@click.option(
    "-v", "--version",
    default=LAST_COMPATIBLE_KONDUIT_VERSION,
    show_default=True,
    help="Only works with the `--download` option to specify the version of konduit-serving to be downloaded."
)
@click.option(
    "-d", "--download",
    is_flag=True,
    help="If set, downloads the pre-built jar file instead of building it. "
         "Also works with `--platform`, `--chip`, `--cuda-version`, `--version` and `--spin` options."
)
def init(platform, https, tag, spin, chip, cuda_version, options, version, download):
    if download:
        if parse(LAST_COMPATIBLE_KONDUIT_VERSION.replace("-SNAPSHOT", "")) < parse(version.replace("-SNAPSHOT", "")):
            print("This version of Python CLI is only compatible with versions <= {}"
                  .format(LAST_COMPATIBLE_KONDUIT_VERSION))
        else:
            download_if_required(get_jar_url(platform, version, spin, chip, cuda_version), KONDUIT_JAR_PATH)
    else:
        git_clone_konduit(https, tag)
        build_jar(platform, spin, chip, cuda_version, options)
