#!/usr/bin/python3
import os
import re
import sys
import tempfile
import time
import socket
import optparse
import traceback

import htcondorce.tools as ce

# HTCondor-CE simple run script.
# Mimics the venerable condor_run

SUBMIT_FILE = """\
universe = grid
grid_resource = condor %(schedd_name)s %(collector_name)s
remote_universe = %(universe)s

executable = %(command)s
arguments = %(arguments)s

output = %(stdout_file)s
error = %(stderr_file)s
log = %(log_file)s

ShouldTransferFiles = YES
WhenToTransferOutput = ON_EXIT
getenv = True

# let condor handle token discovery
use_scitokens = auto
# there's no auto option
use_x509userproxy = %(use_x509userproxy)s
%(extra_attr)s

queue
"""

SUBMIT_REMOTE_FILE = """\
universe = %(universe)s

executable = %(command)s
arguments = %(arguments)s
output = %(stdout_file)s
error = %(stderr_file)s
log = %(log_file)s

ShouldTransferFiles = YES
WhenToTransferOutput = ON_EXIT
getenv = True

# let condor handle token discovery
use_scitokens = auto
# there's no auto option
use_x509userproxy = %(use_x509userproxy)s

+Owner = undefined
+LeaveJobInQueue = ( StageOutFinish > 0 ) =!= true
%(extra_attr)s

queue
"""

def parse_opts():
    parser = optparse.OptionParser()
    parser.add_option("-x", "--extra-attributes-file",
                      help="A file containing additional HTCondor ClassAd attributes to " \
                      "append to the submit ClassAd.",
                      dest="extra_attrs_file")
    parser.add_option("-s", "--schedd-name",
                      help="Name of the schedd to use.",
                      dest="schedd_name")
    parser.add_option("-d", "--debug", help="Print debugging info. Set " \
                      "environmental variable 'CONDOR_CE_RUN_DEBUG' to true to print " \
                      "stacktraces",
                      dest="debug", default=False, action="store_true")
    parser.add_option("-r", "--remote", help="Submit directly to remote " \
                      "schedd, bypassing the local one.", dest="remote", default=False,
                      action="store_true")
    parser.add_option("-l", "--local", help="Have the job run in the CE's " \
                      "local universe.", dest="local", default=False, action="store_true")
    parser.add_option("-n", "--no-clean", help="Do not clean temporary "
                      "files.", dest="clean", default=True, action="store_false")
    parser.add_option("--dry-run", help="Verify arguments and options " \
                      "without submitting a job to the specified CE.", dest="dry_run",
                      default=False, action="store_true")

    parser.disable_interspersed_args()

    opts, args = parser.parse_args()

    return opts, args


def run_command_debug(cmd, debug=False):
    """Run cmd and print stdout/stderr if debug
    """
    rc, stdout, stderr = ce.run_command(cmd)

    if debug:
        print(stdout)
        print(stderr, file=sys.stderr)

    return rc, stdout, stderr


def generate_submit_contents(job_info):
    if job_info['command'][0] != '/':
        for pathdir in os.environ.get("PATH", "").split(":"):
            test_path = os.path.join(pathdir, job_info['command'])
            if os.path.exists(test_path):
                job_info['command'] = test_path
                break

    if job_info['arguments']:
        job_info['arguments'] = '"%s"' % ' '.join(["'%s'" % i for i in \
            job_info['arguments']])
    else:
        job_info['arguments'] = ''

    job_info['use_x509userproxy'] = False
    try:
        if ce.x509_user_proxy_path():
            job_info['use_x509userproxy'] = True
    except OSError:
        pass  # could not open the X.509 user proxy from the expected locations

    if job_info['remote']:
        return SUBMIT_REMOTE_FILE % job_info
    else:
        return SUBMIT_FILE % job_info

def submit_job(job_info, debug=False):
    # TODO: Here and in wait_job below, make sure we can tell the difference
    # between condor_submit failure and failure to exec.
    args = ['condor_submit', job_info['submit_file']]
    if debug:
        args.append("-debug")

    if job_info['remote']:
        args += ['-remote', job_info['schedd_name'], '-pool', job_info['collector_name']]

    rc, stdout, stderr = run_command_debug(args, debug=debug)
    # Example output: 1 job(s) submitted to cluster 14.
    cluster_re = re.compile(r'(\d+) job\(s\) submitted to cluster (\d+)\.')
    try:
        job_info['cluster'] = cluster_re.search(ce.to_str(stdout)).group(2)
    except AttributeError:
        pass

    if job_info['remote'] and ('cluster' not in job_info):
        raise ce.CondorRunException("Could not parse job cluster from " \
                                    "condor_submit output.\nSTDOUT: %s\n STDERR: %s" % (ce.to_str(stdout), ce.to_str(stderr)))

    if rc < 0:
        raise ce.CondorRunException("Failed to submit job; condor_submit " \
                                    "terminated with signal %d.\nSTDOUT: %s\n STDERR: %s" % (-rc, ce.to_str(stdout), ce.to_str(stderr)))
    elif rc > 0:
        raise ce.CondorRunException("Failed to submit job; condor_submit " \
                                    "exited with code %d.\nSTDOUT: %s\n STDERR: %s" % (rc, ce.to_str(stdout), ce.to_str(stderr)))

def wait_for_job(job_info, debug=False):
    if job_info['remote']:
        return wait_for_job_remote(job_info, debug)

    args = ['condor_wait', job_info['log_file']]
    if debug:
        args.append("-debug")

    rc, stdout, stderr = run_command_debug(args, debug=debug)
    if rc < 0:
        raise ce.CondorRunException("Failed to wait for job; condor_wait " \
                                    "terminated with signal %d.\nSTDOUT: %s\n STDERR: %s" % (-rc, stdout, stderr))
    elif rc > 0:
        raise ce.CondorRunException("Failed to wait for job; condor_wait " \
                                    "exited with code %d.\nSTDOUT: %s\n STDERR: %s" % (-rc, stdout, stderr))

def check_remote_status(job_info, debug=False):
    args = ['condor_q', '-format', '%d\n', 'JobStatus', '-pool',
            job_info['collector_name'], '-name', job_info['schedd_name'],
            job_info['cluster']]
    if debug:
        args.append('-debug')

    rc, stdout, stderr = run_command_debug(args, debug=debug)
    job_status = None
    status_re = re.compile(r"(\d+)")
    for line in ce.to_str(stdout):
        try:
            job_status = int(status_re.match(line).group(1))
        except AttributeError:
            pass

    if rc < 0:
        raise ce.CondorRunException("Failed to check job status; condor_q " \
                                    "terminated with signal %d.\nSTDOUT: %s\n STDERR: %s" % (-rc, stdout, stderr))
    elif rc > 0:
        raise ce.CondorRunException("Failed to check job status; condor_q " \
                                    "exited with code %d.\nSTDOUT: %s\n STDERR: %s" % (-rc, stdout, stderr))

    if job_status == None:
        raise ce.CondorRunException("condor_q did not return a valid job status.")

    return job_status

def wait_for_job_remote(job_info, debug=False):
    while check_remote_status(job_info, debug) != 4:
        time.sleep(1)

    args = ['condor_transfer_data', '-pool', job_info['collector_name'],
            '-name', job_info['schedd_name'], job_info['cluster']]
    if debug:
        args.append('-debug')

    rc, stdout, stderr = run_command_debug(args, debug=debug)
    if rc < 0:
        raise ce.CondorRunException("Failed to retrieve job output sandbox; condor_transfer_data " \
                                    "terminated with signal %d.\nSTDOUT: %s\n STDERR: %s" % (-rc, stdout, stderr))
    elif rc > 0:
        raise ce.CondorRunException("Failed to retrieve job output sandbox; condor_transfer_data " \
                                    "exited with code %d.\nSTDOUT: %s\n STDERR: %s" % (-rc, stdout, stderr))

def print_results(job_info):
    for line in open(job_info["stdout_file"], "r").readlines():
        print(line, end='')

    for line in open(job_info["stderr_file"], "r").readlines():
        print(line, end='')


def main():
    opts, args = parse_opts()
    if opts.remote:
        os.environ.setdefault("CONDOR_CONFIG", "/etc/condor-ce/condor_config")
        os.environ.setdefault('_condor_SEC_CLIENT_AUTHENTICATION_METHODS', 'SCITOKENS,SSL,FS')

    if len(args) < 2:
        print("Usage: condor_ce_run <hostname> <command> [arg1] [arg2] [...]")
        return 1

    collector_hostname = args[0].split(":")[0]
    collector_hostname = socket.getfqdn(collector_hostname)
    port_info = args[0].split(":")[1:]
    if port_info:
        collector_hostname = "%s:%s" % (collector_hostname, port_info[0])
    else:
        collector_hostname = "%s:9619" % collector_hostname
    job_info = {'collector_name': collector_hostname, 'command': args[1]}
    if opts.schedd_name:
        job_info['schedd_name'] = opts.schedd_name
    else:
        job_info['schedd_name'] = collector_hostname.split(":")[0]

    os.environ.setdefault("_condor_TOOL_DEBUG", "D_CAT D_ALWAYS:2 D_SECURITY:2")

    job_info['remote'] = opts.remote
    job_info['arguments'] = args[2:]

    if opts.local:
        job_info['universe'] = "local"
    else:
        job_info['universe'] = "vanilla"

    try:
        attr_file = open(opts.extra_attrs_file, 'r')
        job_info['extra_attr'] = attr_file.read()
        attr_file.close()
    except TypeError:
        # No attribute file was specified but the extra_attr key is expected by the submit files
        job_info['extra_attr'] = ''
    except IOError as exc:
        ce.print_formatted_msg(f"Could not open extra attributes file. {exc}")
        sys.exit(1)

    job_info.update(ce.generate_job_files())
    run_script = generate_submit_contents(job_info)

    try:
        if opts.dry_run:
            ce.print_formatted_msg("Without the --dry-run option, condor_ce_run would submit the following " \
                                   "file to %s." % job_info["schedd_name"])
            print(f"\n{run_script}") # don't wrap the submit file
        else:
            pid = os.getpid()
            submit_fd, job_info['submit_file'] = tempfile.mkstemp(dir=".", prefix=".submit_%d_" % pid)
            os.write(submit_fd, ce.to_bytes(run_script))
            os.close(submit_fd)
            submit_job(job_info, opts.debug)
            wait_for_job(job_info, opts.debug)
            print_results(job_info)
    finally:
        if opts.clean or opts.dry_run:
            ce.cleanup_job_files(job_info)

    return 0

if __name__ == '__main__':
    try:
        sys.exit(main())
    except ce.CondorRunException as exc:
        try:
            if os.environ['CONDOR_CE_RUN_DEBUG'].lower() == 'true':
                traceback.print_exc()
        except KeyError:
            ce.print_timestamped_msg(exc)
