#!/usr/bin/env python3

# SPDX-FileCopyrightText: 2009 Fermi Research Alliance, LLC
# SPDX-License-Identifier: Apache-2.0

# Description:
#   This tool creates the configuration file containing
#   the secondary schedds
#

from __future__ import print_function
import sys
import os
import subprocess
import string
import copy


def usage():
    print("This tool creates the 11_gwms_secondary_schedds.config file")
    print("in the HTCondor's config.d directory")
    print()
    print("Usage:")
    print(" glidecondor_createSecSched [-commonlog] [-nocreatedirs] <scheddlist>")
    print("where:")
    print("  -commonlog    - If present, all the schedds will share a common log (default: one log per process)")
    print("  -nocreatedirs - If present, do not (re-)create the needed directories (default: do create the directories)")
    print("  scheddlist    - List of secondary schedds to put in the config file (required)")
    print("Example:")
    print("  glidecondor_createSecSched schedd_glideins1,schedd_glideins2")
    return


def get_config_val(attr, fail_if_missing=True):
    try:
        p = subprocess.Popen(['condor_config_val', attr], stdout=subprocess.PIPE)
    except OSError as e:
        print("Count not find condor_config_val!")
        print("%s\n" % e)
        sys.exit(2)

    rc = p.wait()
    if rc != 0:
        if fail_if_missing:
            print("Attribute '%s' not found" % attr)
            sys.exit(2)
        else:
            return None

    res = p.communicate()
    return res[0].strip('\n')  # only the first line, and no trailing newline


def extract_condor_info():
    global config_dir
    config_dir = get_config_val('LOCAL_CONFIG_DIR')
    global local_dir
    local_dir = get_config_val('LOCAL_DIR')
    global log_dir
    log_dir = get_config_val('LOG')
    global schedlog
    schedlog = get_config_val('SCHEDD_LOG')
    global shadowlog
    shadowlog = get_config_val('SHADOW_LOG')

    global shport_ad
    shport_ad = get_config_val('SHARED_PORT_DAEMON_AD_FILE')
    global shport_sock
    shport_sock = get_config_val('DAEMON_SOCKET_DIR')


def add_schedd(schedd, config_fd):
    # first populate the config file
    attrname = ""
    for c in schedd:
        if c in (string.ascii_uppercase+string.digits):
            attrname += c
        elif c in string.ascii_lowercase:
            attrname += c.upper()
        # drop all others

    env_arr = []
    env_arr.append(("_CONDOR_SCHEDD_NAME", "%(schedd)s"))
    env_arr.append(("_CONDOR_LOCAL_DIR", "%(localdir)s/%(schedd)s"))
    env_arr.append(("_CONDOR_LOCK", "%(localdir)s/%(schedd)s/lock"))
    global common_log
    if common_log:
        env_arr.append(("_CONDOR_SCHEDD_LOG", "%(schedlog)s.%(schedd)s"))
        env_arr.append(("_CONDOR_SHADOW_LOG", "%(shadowlog)s.%(schedd)s"))

    condor_env_arr = []
    for a in env_arr:
        # convert in key=value pair
        # use Condor valriables in the value part
        condor_env_arr.append("%s=%s" % (a[0], a[1] % {'schedd': schedd,
                                                       'localdir': '$(LOCAL_DIR)', 'logdir': '$(LOG)',
                                                       'schedlog': '$(SCHEDD_LOG)', 'shadowlog': '$(SHADOW_LOG)'
                                                       }))

    expenv_str = '%(attrname)s_EXPENV = ' % {'attrname': attrname}
    condor_env_str = " \\\n"+(" "*len(expenv_str)).join(condor_env_arr)

    config_fd.write(("# Secondary schedd %(schedd)s\n" +
                     "%(attrname)s             = $(SCHEDD)\n" +
                     "%(attrname)s_EXPENV = %(envstr)s\n" +
                     '%(attrname)s_ENVIRONMENT = "$(%(attrname)s_EXPENV) $(PRESERVE_SHPORT_EXPENV)"\n' +
                     "SEC_SCHEDD_LIST = $(SEC_SCHEDD_LIST) %(attrname)s\n\n"
                     ) % {'attrname': attrname, 'schedd': schedd, 'envstr': condor_env_str})

    global create_dirs
    if create_dirs:
        # then run condor_init with the modified environment
        config_fd.flush()  # let's just make sure we don't have a patial -> invalid file

        global local_dir, log_dir, schedlog, shadowlog
        # we must start with the old environemnt
        sp_env_dict = copy.deepcopy(os.environ)
        # and just add the new attributes to it
        for a in env_arr:
            # environment is a dictionary
            # use actual paths in the val part
            sp_env_dict[a[0]] = a[1] % {'schedd': schedd,
                                        'localdir': local_dir, 'logdir': log_dir,
                                        'schedlog': schedlog, 'shadowlog': shadowlog}
        # then add the shport part
        global shport_ad, shport_sock
        sp_env_dict['_CONDOR_SHARED_PORT_DAEMON_AD_FILE'] = shport_ad
        sp_env_dict['_CONDOR_DAEMON_SOCKET_DIR'] = shport_sock

        try:
            p = subprocess.Popen(['condor_init'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=sp_env_dict)
        except OSError as e:
            print("Count not find condor_init!")
            print("%s" % e)
            sys.exit(2)

        rc = p.wait()
        res = p.communicate()
        if rc != 0:
            print("condor_init failed for %s: %s" % (schedd, res))
            sys.exit(2)
        print("Schedd %s configured and initialized" % schedd)
    else:
        print("Schedd %s put in config, but not initialized" % schedd)


def create_config(fname, schedds):
    try:
        fd = open(fname, "w")
    except IOError as e:
        print("%s" % e)
        sys.exit(2)

    with fd:
        fd.write("###########################################\n" +
                 "# This file contains the secondary schedds\n" +
                 "# Generated by glidecondor_createSecSched\n" +
                 "# DO NOT modify by hand\n" +
                 "###########################################\n\n" +
                 "PRESERVE_SHPORT_EXPENV= _CONDOR_SHARED_PORT_DAEMON_AD_FILE=$(SHARED_PORT_DAEMON_AD_FILE) \\\n" +
                 "                        _CONDOR_DAEMON_SOCKET_DIR=$(DAEMON_SOCKET_DIR)\n\n")

        for s in schedds:
            add_schedd(s, fd)

        fd.write("DAEMON_LIST = $(DAEMON_LIST) $(SEC_SCHEDD_LIST)\n" +
                 "# we assume we are the only ones re-defining DC_DAEMON_LIST\n" +
                 "DC_DAEMON_LIST = + $(SEC_SCHEDD_LIST)\n")


def parse_args(args):
    global common_log, create_dirs, schedds

    common_log = False
    create_dirs = True

    if len(args) < 1:
        usage()
        sys.exit(1)

    if args[0] == '-h':
        usage()
        sys.exit(0)

    while len(args) > 1:
        if args[0] == '-commonlog':
            common_log = True
        elif args[0] == '-nocreatedirs':
            create_dirs = False
        else:
            print("Unknown option %s" % args[0])
            usage()
            sys.exit(1)
        args = args[1:]

    schedd_list_str = args[0]
    schedds = schedd_list_str.split(',')


def main(args):
    global common_log, create_dirs, schedds
    parse_args(args)

    extract_condor_info()
    global config_dir
    conf_fname = os.path.join(config_dir, "11_gwms_secondary_schedds.config")

    create_config(conf_fname, schedds)


if __name__ == '__main__':
    main(sys.argv[1:])
