#!/bin/sh
"""":
if type python3 >/dev/null 2>&1
then
    exec python3 "$0" "$@"
elif type python >/dev/null 2>&1
then
    exec python "$0" "$@"
else
    echo >&2 "Python is not installed"
fi
exit 1
""" #"
from __future__ import print_function
from contextlib import closing
from logging import handlers
from subprocess import check_output, CalledProcessError
import errno, fileinput, grp, json, logging, os, pwd, re, shlex, shutil

def devnull_init():
    try:
        from subprocess import DEVNULL # py3k
    except ImportError:
        global DEVNULL
        DEVNULL = open(os.devnull, 'wb')

def env_from_file(key):
    try:
        for env in open('/etc/vnlk/vnlk.conf'):
            lst = env.strip().split(' = ', 1)
            if len(lst) == 2 and lst[0] == key:
                return lst[1]
    except:
        pass

VAE = '/opt/vae'
AV = '/opt/avatar'
AV_CONF = '/var/avatar/conf'
VNLK_VAR = '/var/vnlk'
VNLK_CONF = os.path.join(VNLK_VAR, 'conf')
VNLK_VAE = os.path.join(VNLK_VAR, 'vae')
VNLK_VAE_LIB = '/usr/lib/vnlk/vae'
VAE_DISCOVERED = os.path.join(VNLK_VAR, 'vae', '.discovered')
VAE_META_SCHEMAS_DIR = '/tmp/vae-meta-schemas'
#GENERAL_VAE_LAUNCH = 'exec docker run --rm -iv /tmp/$1.$2.sock:/tmp/$1.$2.sock:Z -v {}/$1/conf:{}/$1/conf:ro overcast/vae-$2 {}/$2/bin/vae-engine $@ >/tmp/$1.$2.log 2>&1'.format(APL_CONF, APL_CONF, VAE)
#GENERAL_VAE_LAUNCH_WITH_DEBUG = 'exec docker run --rm -iv /tmp:/tmp:z -v {}/$1/conf:{}/$1/conf:ro --security-opt seccomp=unconfined overcast/vae-$2 gdb -x /tmp/cmds {}/$2/bin/vae-engine $@ >/tmp/$1.$2.log 2>&1'.format(APL_CONF, APL_CONF, VAE)
VAE_ENGINES_IGNORE_LIST = ( \
    'anpr-sdk', 'proxy', 'template', 'validator', 'nodered-config', 'x-server', \
    'anpr', 'crowd', 'crowd2', 'hunt', 'ship', 'torch-ssd', 'weapon')
DATE_FMT = '%Y/%m/%d %H:%M:%S'
LOG_FMT  = '%(asctime)s %(levelname)s - %(message)s'
LOG_PATH = os.path.join(VNLK_VAR, 'log', 'vae-discovery.log')

cap_lst = []
vae_js  = {}

def log_setup():
    global logger
    log_handler = handlers.RotatingFileHandler(LOG_PATH, maxBytes=1024*1024, backupCount=1)
    log_handler.setFormatter(logging.Formatter(LOG_FMT, datefmt = DATE_FMT))
    logger = logging.getLogger(__name__)
    logger.addHandler(log_handler)
    logger.setLevel(logging.INFO)

    #os.chown(LOG_PATH, pwd.getpwnam('apl').pw_uid, grp.getgrnam('apl').gr_gid)

def parseDiscoveryOutput(disco):
    lst = shlex.split(disco)
    obj = json.loads(check_output(lst).decode() if len(lst) > 1 else check_output(disco, shell=True).decode())
    rt  = obj['runtime']
    id  = obj['identity']
    name = id['name']
    vae_js[name] = {'name':name,'minFps':rt['minFPS'],'maxFps':rt['targetFPS'],'title':id.get('title', name.upper()),'version':id.get('version', 'unknown')}

def parseCapabilities(cap_fn, vae_name):
    try:
        with open(cap_fn) as cap_file:
            for line in cap_file:
                line = line.strip()
                lst = line.split('=', 1)
                if len(lst) == 2 and lst[0] == 'STAT_VAE_CAPABILITIES_' + vae_name.upper() and json.loads(lst[1]):
                    cap_lst.append(line)
                    logger.debug(line)
    except:
        logger.exception("While processing capabilities of the 'vae-{}':".format(vae_name))

def printFileContent(fn, level):
    try:
        with open(fn) as f:
            for line in f:
                logger.log(level, line.strip())
    except:
        logger.exception("While print content of the '{}':".format(fn))

def ignore_absent(func, path, exc_inf):
    except_instance = exc_inf[1]
    if isinstance(except_instance, OSError) and \
        except_instance.errno == errno.ENOENT:
        return
    raise except_instance

def modify_start_vae_engine_script(filename):
    with closing(fileinput.FileInput(filename,
                 inplace = True, backup ='.bak')) as f:

        for line in f:
            line = line.replace(":$APL_CONF/av/stat.vae",
                                ":{}/av/stat.vae".format(AV_CONF))
            line = line.replace("$APL_CONF/av/stat.vae",
                                "{}/log/vae.stat".format(VNLK_VAR))
            line = line.replace("-v $APL/etc/cam_mretr.conf:$APL/etc/cam_mretr.conf:ro ",
                                "")
            line = line.replace("-v $APL/etc/md.conf:$APL/etc/md.conf:ro ",
                                "")
            line = line.replace(":$APL_CONF",
                                ":{}".format(AV_CONF))
            line = line.replace("$APL_CONF",
                                VNLK_CONF)
            line = line.replace(":$APL_VAR",
                                ":{}".format(VNLK_VAR))
            line = line.replace("$APL_VAR",
                                VNLK_VAR)
            line = line.replace(":$APL",
                                ":{}/conf".format(AV))
            line = line.replace("$APL/etc",
                                VNLK_CONF)
            print(line, end ='')

def make_executable(path):
    mode = os.stat(path).st_mode
    os.chmod(path, mode | 0o111)

def save_file(fn, content, executable = False):
    try:
        logger.info("Saving '{}'...".format(fn))
        with open(fn, 'w') as f:
            f.write(content)
        f.closed
        if executable:
            make_executable(fn)
    except:
        logger.exception("Fail to write to '{}':".format(fn))

def parseCUDACapableGPUs():
    ret = []
    try:
        with open('/etc/docker/daemon.json') as cfg:
            if json.load(cfg)['runtimes']['nvidia']['path'] is not None:
                gpus = check_output(['nvidia-smi', '-L'], stderr=DEVNULL).decode().split('\n')
                #print(gpus)

                #GPU 0: NVIDIA GeForce GTX 1050 Ti (UUID: GPU-c9b5cc42-a3d5-04b7-d927-c91cbf6924e0)
                pattern = re.compile(r"^GPU\s*(?P<n>\d+):\s*(?P<name>.+)\s+\(UUID:\s*GPU-(?P<uuid>[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})\)$")
                for gpu in gpus:
                    match = pattern.match(gpu)
                    if not match:
                        continue

                    n, name, uuid =          \
                        match.group('n'),    \
                        match.group('name'), \
                        match.group('uuid')

                    ret.append({'name': name, 'uuid': uuid})
    except:
        pass
        #logger.exception("Fail to parse CUDA capable GPUs")
    return ret

def repo_use_latest_tag():
    ret = None
    try:
        with open(os.path.join(VNLK_VAR, 'update', 'settings.json')) as settings:
            ret = json.load(settings)['list']['REGISTRY_USE_LATEST_TAG']
    except:
        pass
    return ret

def in_file(filename, substring):
    """Return True if substring is contained with the named file"""
    ret = False
    try:
        with open(filename, 'r') as fp:
            for line in fp:
                if substring in line:
                    return True
    except:
        pass
    return ret

def clean_dir_content(dir):
    try:
        for f in os.listdir(VAE):
            path = os.path.join(VAE, f)
            try:
                if os.path.isfile(path) or os.path.islink(path):
                    os.unlink(path)
                elif os.path.isdir(path):
                    shutil.rmtree(path, onerror=ignore_absent)
            except Exception as e:
                logger.exception("Failed to delete '{}', reason: {}".format(path, e))
    except:
        pass


if __name__ == '__main__':
    try:
        os.remove(VAE_DISCOVERED)
    except:
        pass

    devnull_init()
    log_setup()
    logger.info('================================================================================')

    # First, let's clean up an old contents of the VAE folder
    clean_dir_content(VAE)

    try:
        if os.path.isdir(VAE_META_SCHEMAS_DIR):
            clean_dir_content(VAE_META_SCHEMAS_DIR)
        else:
            os.makedirs(VAE_META_SCHEMAS_DIR, mode=0o755)
        os.chown(VAE_META_SCHEMAS_DIR, pwd.getpwnam('apl').pw_uid, grp.getgrnam('apl').gr_gid)
    except:
        pass

    # pass through all available docker VAE images
    js = {}
    vae_images_info = check_output('docker images --format "{{.Repository}} {{.ID}} {{.Tag}}"', shell=True).decode().split('\n')
    pattern = re.compile(r"^(?P<repo>.*\/*(overcast|vsaas)\/vae-(?P<name>\S+))\s+(?P<id>\w+)\s+(?P<tag>(?!latest)\S+)$")
    pattern_short_latest = re.compile(r"^(?P<repo>overcast\/vae-(?P<name>\S+))\s+(?P<id>\w+)\s+latest$")
    tmp_dir = os.path.join('/tmp', 'copy-to-host-vae')
    for vae_info in vae_images_info:
        match = pattern.match(vae_info)
        if not match:
            continue

        repo, vae_name, id, tag = \
            match.group('repo'),  \
            match.group('name'),  \
            match.group('id'),    \
            match.group('tag')

        js[vae_name].update({id: {'tag': tag, 'repo': repo}}) if vae_name in js else \
            js.update({vae_name: {id: {'tag': tag, 'repo': repo}}})

    use_latest_tag = repo_use_latest_tag()
    for vae_info in vae_images_info:
        try:
            match = pattern_short_latest.match(vae_info)
            if not match:
                continue

            vae_image, vae_name, id = \
                match.group('repo'),  \
                match.group('name'),  \
                match.group('id')

            vae_dir = os.path.join(VAE, vae_name)
            shutil.rmtree(vae_dir, onerror=ignore_absent)
            shutil.rmtree(tmp_dir, onerror=ignore_absent)
            if vae_name in VAE_ENGINES_IGNORE_LIST:
                continue
            src_dir = os.path.join(vae_dir, 'copy-to-host-vae')
            mschema = os.path.join(vae_dir, 'etc', 'meta-schema.json')
            vae_bin = os.path.join(vae_dir, 'bin')
            disco   = os.path.join(vae_bin, 'discovery')
            run_eng = os.path.join(vae_bin, 'start-vae-engine')
            vae_cap = os.path.join(vae_bin, 'vae-capabilities')
            run     = os.path.join(vae_bin, 'docker.run')
            cap_txt = os.path.join('/tmp', 'vae-capabilities.{}.txt'.format(vae_name))
            tmp_meta_schema_ver = os.path.join('/tmp', 'vae-meta_schema_ver.{}.tmp'.format(vae_name))
            meta_schema_ver = ''
            logger.info("Executing '{}' from '{}' docker image...".format(disco, vae_image))
            cmd = '-v /tmp:/tmp {0} sh -c "[ -d {1} ] && cp -rf {1} /tmp/ >/dev/null 2>&1; [ -x {2} ] && {2} > {3} 2>&1; [ -n $META_SCHEMA_VER ] && echo -n $META_SCHEMA_VER > {4}; {5}"'.format(vae_image, src_dir, vae_cap, cap_txt, tmp_meta_schema_ver, disco)
            parseDiscoveryOutput('docker run --rm ' + cmd)
            if os.path.isfile(tmp_meta_schema_ver):
                meta_schema_ver = check_output('cat ' + tmp_meta_schema_ver, shell=True).decode()
                os.remove(tmp_meta_schema_ver)
            if os.path.isdir(tmp_dir):
                shutil.rmtree(vae_dir, onerror=ignore_absent)
                shutil.move(tmp_dir, vae_dir)
                if os.path.isfile(mschema):
                    if len(meta_schema_ver) != 8: # sanity check;)
                        meta_schema_ver = check_output("cksum {} | cut -d' ' -f1 | xargs printf '%X'".format(mschema), shell=True).decode()
                    logger.info("Meta schema version for 'vae-{}': {}".format(vae_name, meta_schema_ver))
                    with open(mschema) as js_mschema:
                        save_file(os.path.join(VAE_META_SCHEMAS_DIR, '{}-{}.json'.format(vae_name, meta_schema_ver)), \
                                  json.dumps({'engine': vae_name, 'version': meta_schema_ver, 'schema': json.load(js_mschema)}))
            if os.path.isfile(run):
                cmd = '-v /tmp:/tmp {0} sh -c "[ -x {1} ] && {1} > {2} 2>&1" >/dev/null 2>&1'.format(vae_image, vae_cap, cap_txt)
                try:
                    check_output(shlex.split('{} {}'.format(run, cmd)), shell=True).decode()
                except CalledProcessError as e:
                    logger.warning("Command '{} {}' returned non-zero exit status {} while getting capabilities for 'vae-{}'".format(run, cmd, e.returncode, vae_name))
                    logger.log(logging.WARNING, "'{}' content:".format(cap_txt))
                    printFileContent(cap_txt, logging.WARNING)
                    os.remove(cap_txt)
                except:
                    logger.exception("While getting capabilities for 'vae-{}'".format(vae_name))
            if os.path.isfile(cap_txt):
                logger.info("Trying to parse capabilities for 'vae-{}'...".format(vae_name))
                parseCapabilities(cap_txt, vae_name)
                os.remove(cap_txt)
            if os.path.isfile(run_eng):
                modify_start_vae_engine_script(run_eng)

            # let's fill VAE version
            ver = js[vae_name][id]['tag'] \
                if vae_name in js and id in js[vae_name] else 'unknown'
            vae_js[vae_name]['version'] = 'unknown' \
                if use_latest_tag != '1' and ver == 'latest' else ver
        except:
            shutil.rmtree(tmp_dir, onerror=ignore_absent)
            logger.exception("While processing '{}' docker image:".format(vae_image))

    # let's pass through all local VAE engines (which are part of Link)
    for eng in os.listdir(VNLK_VAE):
        vae_dir = os.path.join(VNLK_VAE, eng)
        disco   = os.path.join(vae_dir, 'bin', 'discovery')
        mschema = os.path.join(vae_dir, 'etc', 'meta-schema.json')
        dst_dir = os.path.join(VAE, eng)
        mschema_ver = ''
        try:
            if eng in VAE_ENGINES_IGNORE_LIST \
               or not os.access(os.path.join(VNLK_VAE_LIB, 'libvae-{}.so'.format(eng)), os.F_OK) \
               or not os.access(disco, os.F_OK):
               #or not in_file(disco, '"shared-library"'):
                continue

            shutil.rmtree(dst_dir, onerror=ignore_absent)
            if eng in vae_js:
                del vae_js[eng]

            logger.info("Executing '{}' ...".format(disco))
            parseDiscoveryOutput(disco)
            shutil.copytree(vae_dir, dst_dir)
            if os.path.isfile(mschema):
                mschema_ver = check_output("cksum {} | cut -d' ' -f1 | xargs printf '%X'".format(mschema), shell=True).decode()
                if len(meta_schema_ver) == 8: # sanity check;)
                    logger.info("Meta schema version for 'vae-{}': {}".format(eng, mschema_ver))
                    with open(mschema) as js_mschema:
                        save_file(os.path.join(VAE_META_SCHEMAS_DIR, '{}-{}.json'.format(eng, mschema_ver)), \
                                  json.dumps({'engine': eng, 'version': mschema_ver, 'schema': json.load(js_mschema)}))
        except Exception as e:
            logger.exception("While processing local VAE '{}', reason: {}".format(eng, e))
            shutil.rmtree(dst_dir, onerror=ignore_absent)
            if eng in vae_js:
                del vae_js[eng]

    # Write STAT_VAE_LIST
    output = ''
    for vae_name in vae_js:
        item = vae_js[vae_name]
        output += ',' + json.dumps(item) if output else json.dumps(item)

    output = 'STAT_VAE_LIST=[' + output + ']'
    print(output)
    cuda_gpus = parseCUDACapableGPUs()
    #print(cuda_gpus)

    logger.info(output)
    cap_lst.append(output)
    if len(cuda_gpus) != 0:
        output = 'STAT_VAE_CUDA_CAPABLE_GPUS=' + json.dumps(cuda_gpus)
        logger.info(output)
        cap_lst.append(output)
    save_file(os.path.join(VNLK_VAR, 'log', 'vae.stat'), '[general]\n' + '\n'.join(cap_lst) + '\n')

    # write general start script for docker containers w/o such 'launchScript' attribute in discovery
    #fn = os.path.join(APL_VAR, 'vae', 'start-vae-engine')
    #if not os.access(fn, os.F_OK):
    #    save_file(fn, '#!/bin/sh\n\n{}\n#{}\n'.format(GENERAL_VAE_LAUNCH, GENERAL_VAE_LAUNCH_WITH_DEBUG), True)

    try:
        with open(VAE_DISCOVERED, 'w') as f:
            os.chown(VAE_DISCOVERED, pwd.getpwnam('apl').pw_uid, grp.getgrnam('apl').gr_gid)
    except:
        pass
