#!/usr/bin/python3

import os
import re
import pwd
import classad2 as classad
from collections import OrderedDict

JOB_ROUTER_CONFIG = r"""JOB_ROUTER_TRANSFORM_Env @=jrt
    EVALMACRO default_env {advertise_pilots}
    if $(USE_CE_HOME_DIR)
        EVALMACRO use_ce_home userHome(Owner, "/")
        default_env = $(default_env) HOME=$(use_ce_home)
    endif

    if ! defined default_pilot_job_env
        default_pilot_job_env = ""
    endif

    SET osg_environment "{osg_environment}"
    EVALSET environment mergeEnvironment("$(default_env)", \
                                         osg_environment, \
                                         orig_environment, \
                                         $(CONDORCE_PILOT_JOB_ENV), \
                                         $(default_pilot_job_env))
@jrt

JOB_ROUTER_POST_ROUTE_TRANSFORM_NAMES = $(JOB_ROUTER_POST_ROUTE_TRANSFORM_NAMES) Env

JOB_ROUTER_DEFAULTS_GENERATED @=jrd
  [ MaxIdleJobs = 2000;
    MaxJobs = $(CONDORCE_MAX_JOBS);
    /* by default, accept all jobs */
    Requirements = True;

    /* now modify routed job attributes */
    /* remove routed job if the client disappears for 48 hours or it is idle for 6 */
    /*set_PeriodicRemove = (LastClientContact - time() > 48*60*60) ||
                           (JobStatus == 1 && (time() - QDate) > 6*60);*/
    delete_PeriodicRemove = true;
    delete_CondorCE = true;
    delete_TotalSubmitProcs = true;
    set_RoutedJob = true;

    /* Copy AuthToken attributes if they exist - the routed job will have the original attributes deleted */
    copy_AuthTokenSubject = "orig_AuthTokenSubject";
    copy_AuthTokenIssuer = "orig_AuthTokenIssuer";
    copy_AuthTokenGroups = "orig_AuthTokenGroups";
    copy_AuthTokenScopes = "orig_AuthTokenScopes";
    copy_AuthTokenId = "orig_AuthTokenId";

    /* Set the environment */
    copy_environment = "orig_environment";
    set_osg_environment = "{osg_environment}";
    eval_set_environment = mergeEnvironment(join(" ",
                                                 $(USE_CE_HOME_DIR) =?= True ?
                                                     strcat("HOME=", userHome(Owner, "/")) :
                                                     "",
                                                 {advertise_pilots}),
                                            osg_environment,
                                            orig_environment,
                                            $(CONDORCE_PILOT_JOB_ENV),
                                            default_pilot_job_env);

    /* Set new requirements */
    /* set_requirements = LastClientContact - time() < 30*60; */
    set_requirements = (RequestGpus?:0) >= (TARGET.Gpus?:0);

    /* Note default memory request of 2GB */
    /* Note yet another nested condition allow pass attributes (maxMemory,xcount,jobtype,queue)
       via gWMS Factory described within ClassAd */
    eval_set_OriginalMemory = ifThenElse(maxMemory isnt undefined,
                                         maxMemory,
                                         ifThenElse(default_maxMemory isnt undefined,
                                                    default_maxMemory,
                                                    2000));

    /* Duplicate OriginalMemory expression and add remote_ prefix.
       This passes the attribute from gridmanager to BLAHP. */
    eval_set_remote_OriginalMemory = ifThenElse(maxMemory isnt undefined,
                                                maxMemory,
                                                ifThenElse(default_maxMemory isnt undefined,
                                                           default_maxMemory,
                                                           2000));
    set_JOB_GLIDEIN_Memory = "$$(TotalMemory:0)";
    set_JobMemory = JobIsRunning ? int(MATCH_EXP_JOB_GLIDEIN_Memory)*95/100 : OriginalMemory;

    set_RequestMemory = ifThenElse(WantWholeNode is true,
                                   !isUndefined(TotalMemory) ? TotalMemory*95/100 : JobMemory,
                                   OriginalMemory);
    eval_set_remote_queue = ifThenElse(batch_queue isnt undefined,
                            batch_queue,
                            ifThenElse(queue isnt undefined,
                                       queue,
                                       ifThenElse(default_queue isnt undefined,
                                                  default_queue,
                                                  "")));

    /* Request GPUs for whole node jobs (HTCONDOR-103) */
    /* If a whole node job requests GPUs and is matched to a machine with GPUs then set the job's RequestGPUs to all the GPUs on that machine */
    copy_RequestGPUs = "orig_RequestGPUs";
    copy_WantWholeNode = "WholeNodes";
    eval_set_OriginalGPUs = orig_RequestGPUs;
    /* MATCH_EXP_JOB_GLIDEIN_GPUs will be based on JOB_GLIDEIN_GPUs (set below) once the routed job is matched to an HTCondor slot */
    set_GlideinGPUsIsGood = !isUndefined(MATCH_EXP_JOB_GLIDEIN_GPUs) && (int(MATCH_EXP_JOB_GLIDEIN_GPUs) isnt error);
    /* JobGPUs set below; TotalGPUs comes from the slot ad, WantWholeNode from the job ad */
    set_JOB_GLIDEIN_GPUs = "$$(ifThenElse(WantWholeNode is true, !isUndefined(TotalGPUs) ? TotalGPUs : JobGPUs, OriginalGPUs))";
    set_JobGPUs = JobIsRunning ? int(MATCH_EXP_JOB_GLIDEIN_GPUs) : OriginalGPUs;
    set_RequestGPUs = ifThenElse((WantWholeNode is true && OriginalGPUs isnt undefined),
                                 (!isUndefined(TotalGPUs) && TotalGPUs > 0)? TotalGPUs : JobGPUs,
                                 OriginalGPUs);


    /* HTCondor uses RequestCpus; blahp uses SMPGranularity and NodeNumber.  Default is 1 core. */
    copy_RequestCpus = "orig_RequestCpus";
    eval_set_OriginalCpus = ifThenElse(xcount isnt undefined,
                                       xcount,
                                       ifThenElse(orig_RequestCpus isnt undefined,
                                                  ifThenElse(orig_RequestCpus > 1,
                                                             orig_RequestCpus,
                                                             ifThenElse(default_xcount isnt undefined,
                                                                        default_xcount,
                                                                        1)),
                                                  ifThenElse(default_xcount isnt undefined,
                                                             default_xcount,
                                                             1)));
    set_GlideinCpusIsGood = !isUndefined(MATCH_EXP_JOB_GLIDEIN_Cpus) && (int(MATCH_EXP_JOB_GLIDEIN_Cpus) isnt error);
    set_JOB_GLIDEIN_Cpus = "$$(ifThenElse(WantWholeNode is true, !isUndefined(TotalCpus) ? TotalCpus : JobCpus, OriginalCpus))";
    set_JobIsRunning = (JobStatus =!= 1) && (JobStatus =!= 5) && GlideinCpusIsGood;
    set_JobCpus = JobIsRunning ? int(MATCH_EXP_JOB_GLIDEIN_Cpus) : OriginalCpus;
    set_RequestCpus = ifThenElse(WantWholeNode is true,
                                 !isUndefined(TotalCpus) ? TotalCpus : JobCpus,
                                 OriginalCpus);
    eval_set_remote_SMPGranularity = ifThenElse(xcount isnt undefined,
                                                xcount,
                                                ifThenElse(default_xcount isnt undefined,
                                                           default_xcount,
                                                           1));
    eval_set_remote_NodeNumber = ifThenElse(xcount isnt undefined,
                                            xcount,
                                            ifThenElse(default_xcount isnt undefined,
                                                       default_xcount,
                                                       1));

    /* BatchRuntime is in seconds but users configure default_maxWallTime and ROUTED_JOB_MAX_TIME in minutes */
    copy_BatchRuntime = "orig_BatchRuntime";
    eval_set_BatchRuntime = ifThenElse(maxWallTime isnt undefined,
                                       60*maxWallTime,
                                       ifThenElse(orig_BatchRuntime isnt undefined,
                                                  orig_BatchRuntime,
                                                  ifThenElse(default_maxWallTime isnt undefined,
                                                             60*default_maxWallTime,
                                                             60*$(ROUTED_JOB_MAX_TIME))));

    set_CondorCE = 1;
    eval_set_CERequirements = ifThenElse(default_CERequirements isnt undefined,
                                         strcat(default_CERequirements, ",CondorCE"),
                                         "CondorCE");

    copy_OnExitHold = "orig_OnExitHold";
    set_OnExitHold = ifThenElse(orig_OnExitHold isnt undefined,
                                orig_OnExitHold,
                                false) ||
                     ifThenElse(minWalltime isnt undefined && RemoteWallClockTime isnt undefined,
                                RemoteWallClockTime < 60*minWallTime,
                                false);
    copy_OnExitHoldReason = "orig_OnExitHoldReason";
    set_OnExitHoldReason = ifThenElse((orig_OnExitHold isnt undefined) && orig_OnExitHold,
                                      ifThenElse(orig_OnExitHoldReason isnt undefined,
                                                 orig_OnExitHoldReason,
                                                 strcat("The on_exit_hold expression (",
                                                        unparse(orig_OnExitHold),
                                                        ") evaluated to TRUE.")),
                                      ifThenElse(minWalltime isnt undefined &&
                                                   RemoteWallClockTime isnt undefined &&
                                                   (RemoteWallClockTime < 60*minWallTime),
                                                 strcat("The job's wall clock time, ",
                                                        int(RemoteWallClockTime/60),
                                                        "min, is less than the minimum specified by the job (",
                                                        minWalltime,
                                                        ")"),
                                                 "Job held for unknown reason."));

    copy_OnExitHoldSubCode = "orig_OnExitHoldSubCode";
    set_OnExitHoldSubCode = ifThenElse((orig_OnExitHold isnt undefined) && orig_OnExitHold,
                                       ifThenElse(orig_OnExitHoldSubCode isnt undefined,
                                                   orig_OnExitHoldSubCode,
                                                   1),
                                       42);
  ]
  @jrd
"""

ACCOUNTING_GROUP_CONFIG = """
# Also append the accounting group routing rules to the old job router syntax
# since it's still the default syntax
JOB_ROUTER_DEFAULTS_GENERATED @=jrd
  $(JOB_ROUTER_DEFAULTS_GENERATED)
  [
    set_AccountingGroupOSG = {accounting_group};
    eval_set_AccountingGroup = AccountingGroupOSG;
  ]
@jrd
"""

osg_environment_files = ["/var/lib/osg/osg-job-environment.conf",
                         "/var/lib/osg/osg-local-job-environment.conf"]


def advertise_pilots_value():
    """Allow admins to prevent pilots from advertising back to the site CE collector
    """
    try:
        if os.environ['DISABLE_GLIDEIN_ADS'].lower() == 'true':
            return '""'
    except KeyError:
        pass

    return 'ifThenElse($(DISABLE_PILOT_ADS) =?= True, "", strcat("CONDORCE_COLLECTOR_HOST=", "$(COLLECTOR_HOST)"))'


def parse_extattr(extattr_file):
    if not os.path.exists(extattr_file):
        return []
    fd = open(extattr_file)
    results = []
    for line in fd:
        if line.startswith("#"):
            continue
        line = line.strip()
        if not line:
            continue
        info = line.rsplit(" ", 1)
        if len(info) != 2:
            continue
        results.append((str(info[0].strip()), str(info[1]).strip()))
    return results


def parse_uids(uid_file):
    if not os.path.exists(uid_file):
        return []
    fd = open(uid_file)
    results = []
    for line in fd:
        if line.startswith("#"):
            continue
        line = line.strip()
        if not line:
            continue
        info = line.split(" ", 1)
        if len(info) != 2:
            continue
        try:
            uid = int(info[0])
            results.append((pwd.getpwuid(uid).pw_name, str(info[1]).strip()))
        except ValueError:
            results.append((info[0], str(info[1]).strip()))
    return results


def set_accounting_group(uid_file="/etc/osg/uid_table.txt", extattr_file="/etc/osg/extattr_table.txt"):
    attr_mappings = parse_extattr(extattr_file)
    uid_mappings = parse_uids(uid_file)
    if not classad:
        return "Owner"
    elif not attr_mappings and not uid_mappings:
        return None
    accounting_group_str = ''
    for mapping in uid_mappings:
        accounting_group_str += 'ifThenElse(Owner == %s, strcat(%s, ".", Owner), ' % (classad.quote(mapping[0]),
                                                                                      classad.quote(mapping[1]))
    for mapping in attr_mappings:
        accounting_group_str += 'ifThenElse(regexp(%s, x509UserProxySubject), strcat(%s, ".", Owner), ' \
                                % (classad.quote(mapping[0]), classad.quote(mapping[1]))
        accounting_group_str += 'ifThenElse(x509UserProxyFirstFQAN isnt Undefined && ' + \
                                'regexp(%s, x509UserProxyFirstFQAN), strcat(%s, ".", Owner), ' \
                                % (classad.quote(mapping[0]), classad.quote(mapping[1]))
    accounting_group_str += "Owner" + ")"*(len(attr_mappings)*2+len(uid_mappings))
    return accounting_group_str


def condor_env_escape(val):
    """
    Escape the environment variable to match Condor's escape sequence.

    From condor_submit's man page:
    1 Put double quote marks around the entire argument string. Any literal
      double quote marks within the string must be escaped by repeating the
      double quote mark.
    2 Use white space (space or tab characters) to separate environment
      entries.
    3 To put any white space in an environment entry, surround the space and
      as much of the surrounding entry as desired with single quote marks.
    4 To insert a literal single quote mark, repeat the single quote mark
      anywhere inside of a section surrounded by single quote marks.

    THIS IS NOT A GENERIC ESCAPER; we assume this only works on the OSG
    environment file format.  We also assume the input is valid.
    """
    if val.startswith('"') and val.endswith('"'):
        val = val[1:-1]
    val = val.replace('\\', '')  # Nuke escape sequences.
    val = val.replace('"', '""')
    val = val.replace("'", "''")
    return "'" + val + "'"


export_line_re = re.compile(r'^export\s+([a-zA-Z_]\w*)')
variable_line_re = re.compile(r'([a-zA-Z_]\w*)=(.+)')
shell_var_re = re.compile(r'"?\$(\w*)"?')


def read_osg_environment_file(filename):
    """
    Parse the OSG environment file.

    This file is maddening because it APPEARS to be a file you can source
    with bash; however, it has a very limited syntax.
    """
    fd = open(filename, 'r')
    export_lines = []
    env = {}
    for line in fd.readlines():
        line = line.strip()
        # Ignore comments
        if line.startswith("#"):
            continue
        m = export_line_re.match(line)
        if m:
            export_lines.append(m.group(1))
        m = variable_line_re.match(line)
        if m:
            (job_var, value) = m.groups()
            shell_var = shell_var_re.match(value)
            if shell_var:
                ce_var = os.getenv(shell_var.group(1))
                if ce_var:
                    env[job_var] = condor_env_escape(ce_var)
            else:
                env[job_var] = condor_env_escape(value)
    return dict([(i[0], i[1]) for i in env.items() if i[0] in export_lines])


def main():
    # Read environment files, preferring the local env
    osg_env = OrderedDict()
    for filename in osg_environment_files:
        try:
            osg_env.update(OrderedDict(read_osg_environment_file(filename)))
        except IOError:
            pass
    # Construct HTCondor-formatted environment string
    env_string = " ".join(["%s=%s" % (i[0], i[1]) for i in osg_env.items()])

    advertise_pilots = advertise_pilots_value()
    defaults = JOB_ROUTER_CONFIG.format(osg_environment=env_string, advertise_pilots=advertise_pilots)

    accounting_group_expr = set_accounting_group()
    if accounting_group_expr:
        defaults += ACCOUNTING_GROUP_CONFIG.format(accounting_group=accounting_group_expr)

    print(defaults)


if __name__ == "__main__":
    main()
