#!/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
#


import copy
import os
import string
import subprocess
import sys


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 = f"{attrname}_EXPENV = "
    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 partial -> invalid file

        global local_dir, log_dir, schedlog, shadowlog
        # we must start with the old environment
        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(f"condor_init failed for {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 OSError 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:])
