#!/usr/bin/python3

import os
import re
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
"""

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 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)

    print(defaults)


if __name__ == "__main__":
    main()
