#!/bin/sh
# coding=utf-8

''':'
if [ "$1" = "init" ]; then
    printf "{\"status\":\"Success\",\"capabilities\":{\"attach\":false}}" >&1
    exit 0
fi
PATH=$PATH:/usr/local/bin
PYTHON=`command -v python2`
if [ -z $PYTHON ]; then
    PYTHON=`command -v python`
fi
if [ -z $PYTHON ]; then
    PYTHON=`command -v python3`
fi
if [ -z $PYTHON ]; then
    echo "no Python available, please install it first"
    exit 1
fi
sudo -E $PYTHON "$0" "$@"
exit $?
'''  # '''

import glob
import os
import sys

try:
    from urllib.parse import urlparse
    from urllib.request import urlopen
except ImportError:
    from urlparse import urlparse
    from urllib2 import urlopen


def log(msg):
    sys.stderr.write(str(msg) + "\n")


class HadoopDist(object):
    def __init__(self):
        self.download_dir = "/usr/local/lib"
        self.jar_link_path = os.path.join(self.download_dir, "juicefs-hadoop.jar")
        self.create_dir_if_not_exist = True
        self.link_dirs = []
        self.java_home = "/usr/lib/jvm/java"
        self.cache_dir = None
        log("setup JuiceFS in %s" % self.__class__.__name__)

    def boot(self, args):
        conf_dir = args.conf_dir
        conf_file = args.conf_file
        cache_dir = self.cache_dir if args.cache_dir is None else args.cache_dir

        # download jar
        self.download_jar(args)

        # symlink
        self._ln_jars()

        # juicefs conf dir
        if conf_dir and not os.path.exists(conf_dir):
            os.makedirs(conf_dir)
            os.chmod(conf_dir, 0o777)
        if conf_file:
            self.download_conf_file(args)

        # juicefs cache dir
        if cache_dir:
            self._make_cache_dir(cache_dir)

    def download_jar(self, args):
        jar = args.jar

        if jar is None:
            jar = "https://s.juicefs.com/static/Linux/juicefs-hadoop-4.9.0.jar"
        if "s.juicefs.com" in jar:
            res = os.system("curl -o %s %s" % (os.path.join(self.download_dir, os.path.basename(jar)), jar))
        else:
            res = self.do_download_jar(args)
        if res != 0:
            log("download from %s failed" % jar)
            sys.exit(1)

        # make symlink
        for p in glob.glob(os.path.join(self.download_dir, "juicefs-hadoop*")):
            if not os.path.exists(self.jar_link_path):
                os.symlink(p, self.jar_link_path)
            os.chmod(p, 0o644)
            break

    def do_download_jar(self, args):
        pass

    def download_conf_file(self, args):
        if self.do_download_conf_file(args) != 0:
            log("download from %s failed" % args.conf_file)
            sys.exit(1)
        for p in glob.glob(os.path.join(args.conf_dir, "*")):
            os.chown(p, 0o644)

    def do_download_conf_file(self, args):
        pass

    def _ln_jars(self):
        for d in self.link_dirs:
            if not os.path.exists(d):
                if self.create_dir_if_not_exist:
                    os.makedirs(d)
                else:
                    continue
            for p in glob.glob(os.path.join(d, "juicefs-hadoop*")):
                os.remove(p)
            jar_link = os.path.join(d, "juicefs-hadoop.jar")
            tools_link = os.path.join(d, "tools.jar")
            os.symlink(self.jar_link_path, jar_link)
            log("install juicefs-hadoop.jar to %s" % jar_link)
            if os.path.islink(tools_link) or os.path.exists(tools_link):
                os.remove(tools_link)
            os.symlink(
                os.path.join(self.java_home, "lib/tools.jar"),
                os.path.join(d, "tools.jar"),
            )

    def _make_cache_dir(self, cache_dir=None):
        patterns = cache_dir.split(':')
        for pattern in patterns:
            dirs = self._expand_dir(pattern)
            for dir in dirs:
                if not os.path.exists(dir):
                    os.makedirs(dir)
                    os.chmod(dir, 0o777)
                    log("make cache dir %s" % dir)

    def _expand_dir(self, pattern):
        chars = '*?['
        if not any(elem in chars for elem in pattern):
            return [pattern]
        parent, name = os.path.split(pattern)
        rest_name = ""
        while not any(elem in chars for elem in name):
            rest_name = os.path.join(name, rest_name)
            parent, name = os.path.split(parent)

        matched = glob.glob(os.path.join(parent, name))
        return [os.path.join(d, rest_name) for d in matched]


class AmazonEMR(HadoopDist):
    def __init__(self):
        super(AmazonEMR, self).__init__()
        self.link_dirs = [
            "/usr/lib/hadoop/lib/",
            "/usr/lib/spark/jars/",
            "/usr/lib/hive/auxlib/",
            "/usr/lib/presto/plugin/hive-hadoop2",
        ]
        self.cache_dir = "/mnt*/jfs"

    def do_download_jar(self, args):
        return os.system("aws s3 cp %s %s" % (args.jar, self.download_dir))

    def do_download_conf_file(self, args):
        conf_dir = args.conf_dir
        d = os.path.join(conf_dir, ".juicefs")
        if not os.path.exists(d):
            os.mkdir(d)
            os.chmod(d, 0o777)
        return os.system("aws s3 cp %s %s" % (args.conf_file, d))


class AliyunEMR(HadoopDist):
    def __init__(self):
        super(AliyunEMR, self).__init__()
        dir_patterns = [
            "/opt/apps/ecm/service/hadoop/*/package/hadoop*/share/hadoop/common/lib",
            "/opt/apps/ecm/service/spark/*/package/spark*/jars",
            "/opt/apps/ecm/service/presto/*/package/presto*/plugin/hive-hadoop2",
            "/opt/apps/ecm/service/impala/*/package/impala*/lib",
        ]

        # remove juicefs-hadoop.jar under presto lib dir
        presto_lib_jfs_jar = glob.glob('/opt/apps/ecm/service/presto/*/package/presto*/lib/juice*')
        if len(presto_lib_jfs_jar) > 0:
            for p in presto_lib_jfs_jar:
                os.remove(p)

        # auto create hive auxlib
        hive_home = glob.glob('/opt/apps/ecm/service/hive/*/package/apache-hive*')
        if len(hive_home) > 0:
            hive_aux_dir = os.path.join(hive_home[0], "auxlib")
            if not os.path.exists(hive_aux_dir):
                os.mkdir(hive_aux_dir)
                os.chmod(hive_aux_dir, 0o775)

        self.java_home = "/usr/lib/jvm/java-openjdk"
        self.link_dirs = [glob.glob(p)[0] for p in dir_patterns if len(glob.glob(p)) > 0]
        self.cache_dir = "/mnt/disk*/jfs"

    def _download_from_oss(self, download_dir, url, endpoint):
        pr = urlparse(url)
        netloc = pr.netloc
        bucket = pr.netloc
        filepath = pr.path if not pr.path.startswith("/") else pr.path[1:]
        if len(netloc.split(".")) > 1:
            endpoint = netloc[netloc.index(".") + 1 :]
            bucket = netloc[: netloc.index(".")]
        file_url = "oss://" + bucket + "/" + filepath
        cmd = "ossutil cp -f "
        try:
            ak = urlopen("http://localhost:10011/role-access-key-id").read()
            sk = urlopen("http://localhost:10011/role-access-key-secret").read()
            token = urlopen("http://localhost:10011/role-security-token").read()
        except:
            import json

            content = json.loads(
                urlopen(
                    "http://100.100.100.200/latest/meta-data/ram/security-credentials/AliyunECSInstanceForEMRRole"
                ).read()
            )
            ak = content["AccessKeyId"]
            sk = content["AccessKeySecret"]
            token = content["SecurityToken"]
        cmd += "-t %s " % token
        cmd += "-i %s " % ak
        cmd += "-k %s " % sk
        if endpoint is not None:
            cmd += "-e %s " % endpoint
        return os.system("%s %s %s" % (cmd, file_url, download_dir))

    def do_download_jar(self, args):
        return self._download_from_oss(self.download_dir, args.jar, args.endpoint)

    def do_download_conf_file(self, args):
        conf_dir = args.conf_dir
        d = os.path.join(conf_dir, ".juicefs")
        if not os.path.exists(d):
            os.mkdir(d)
            os.chmod(d, 0o777)
        return self._download_from_oss(d, args.conf_file, args.endpoint)


class TencentEMR(HadoopDist):
    def __init__(self):
        super(TencentEMR, self).__init__()
        self.create_dir_if_not_exist = False
        self.link_dirs = [
            "/usr/local/service/presto/plugin/hive-hadoop2",
            "/usr/local/service/hadoop/share/hadoop/common/lib",
            "/usr/local/service/spark/jars",
            "/usr/local/service/hive/auxlib",
        ]
        self.java_home = "/usr/local/jdk"
        self.cache_dir = "/data*/jfs"

    def do_download_jar(self, args):
        jar_url = args.jar
        if ".cos." in jar_url and jar_url.startswith("http"):
            jar_url = "cosn://%s%s" % (
                urlparse(jar_url).hostname.split(".")[0],
                urlparse(jar_url).path,
            )
        return os.system("/usr/local/service/hadoop/bin/hadoop fs -get -f %s %s" % (jar_url, self.download_dir))

    def do_download_conf_file(self, args):
        conf_dir = args.conf_dir
        conf_file = args.conf_file
        d = os.path.join(conf_dir, ".juicefs")
        if not os.path.exists(d):
            os.mkdir(d)
            os.chmod(d, 0o777)
        if ".cos." in conf_file:
            conf_file = "cosn://%s%s" % (
                urlparse(conf_file).hostname.split(".")[0],
                urlparse(conf_file).path,
            )
        return os.system("/usr/local/service/hadoop/bin/hadoop fs -get -f %s %s" % (conf_file, d))


class UHadoop(HadoopDist):
    def __init__(self):
        super(UHadoop, self).__init__()
        self.home = "/home/hadoop"
        self.link_dirs = [
            os.path.join(self.home, app)
            for app in [
                "presto/plugin/hive-hadoop2",
                "share/hadoop/common/lib",
                "spark/jars",
                "hive/auxlib",
            ]
        ]


def find_hadoop_dist():
    if os.path.exists("/usr/local/service/hadoop"):
        return TencentEMR()
    if os.path.exists("/usr/bin/aws"):
        return AmazonEMR()
    if os.path.exists("/home/hadoop/share/hadoop"):
        return UHadoop()
    if os.path.exists("/opt/apps/ecm/service/hadoop"):
        return AliyunEMR()
    else:
        return HadoopDist()


def main():
    import argparse

    parser = argparse.ArgumentParser(description='for emr bootstrap')
    parser.add_argument("--cache-dir", dest="cache_dir", help="dir used for cache data in local disk")
    parser.add_argument(
        "--conf-dir",
        dest="conf_dir",
        help="dir used for download JuiceFS configuration",
        default="/etc/juicefs",
    )
    parser.add_argument(
        "-j",
        "--jar",
        help="juicefs hadoop jar download address, can be oss, s3 or s.juicefs.com",
    )
    parser.add_argument("-e", "--endpoint", dest="endpoint", help="endpoint", default=None)
    parser.add_argument("-a", "--access_key", dest="ak", help="secret_id or access_key", default=None)
    parser.add_argument("-s", "--secret_key", dest="sk", help="secret_key", default=None)
    parser.add_argument(
        "-c",
        "--conf-file",
        dest="conf_file",
        help="juicefs conf file, you can find it under " "/root/.juicefs on your mount jfs machine",
        default=None,
    )

    args = parser.parse_args()
    dist = find_hadoop_dist()

    return getattr(dist, "boot")(args)


sys.exit(main() or 0)
