#!/usr/bin/env python3

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

"""This tool prints out the matches between jobs and Factory entries

Arguments:
   $1 = work dir
   $2 = group_name
"""


import os
import os.path
import sys
import traceback

from glideinwms.frontend import glideinFrontendConfig, glideinFrontendInterface, glideinFrontendLib

# /var/lib/gwms-frontend/vofrontend/frontend.descript has hardcoded line
# MatchExpr (True) and (getGlideinCpusNum(glidein) >=
# int(job.get("RequestCpus", 1)))
# and this script fails to find getGlideinCpusNum() unless we import it
# from glideinFrontendLib explicitly
from glideinwms.frontend.glideinFrontendLib import getGlideinCpusNum  # noqa: F401


############################################################
def do_work(elementDescript, paramsDescript, signatureDescript):
    match_obj = elementDescript.merged_data["MatchExprCompiledObj"]

    # query condor

    glidein_dict = {}
    factory_constraint = elementDescript.merged_data["FactoryQueryExpr"]
    factory_pools = elementDescript.merged_data["FactoryCollectors"]
    for factory_pool in factory_pools:
        factory_pool_node = factory_pool[0]
        factory_identity = factory_pool[1]
        my_identity_at_factory_pool = factory_pool[2]
        try:
            full_constraint = '(%s) && ((PubKeyType=?="RSA") && (GlideinAllowx509_Proxy=!=False))' % factory_constraint
            factory_glidein_dict = glideinFrontendInterface.findGlideins(
                factory_pool_node, None, signatureDescript.signature_type, full_constraint
            )
        except RuntimeError as e:
            if factory_pool_node is not None:
                print(f"Failed to talk to factory_pool {factory_pool_node}: {e}")
            else:
                print("Failed to talk to factory_pool: %s" % e)
            # failed to talk, like empty... maybe the next factory will have something
            factory_glidein_dict = {}

        for glidename in factory_glidein_dict.keys():
            if ("AuthenticatedIdentity" not in factory_glidein_dict[glidename]["attrs"]) or (
                factory_glidein_dict[glidename]["attrs"]["AuthenticatedIdentity"] != factory_identity
            ):
                print(f"Found an untrusted factory {glidename} at {factory_pool_node}; ignoring.")
                if "AuthenticatedIdentity" in factory_glidein_dict[glidename]["attrs"]:
                    print(
                        "Found an untrusted factory %s at %s; identity mismatch '%s'!='%s'"
                        % (
                            glidename,
                            factory_pool_node,
                            factory_glidein_dict[glidename]["attrs"]["AuthenticatedIdentity"],
                            factory_identity,
                        )
                    )
            else:
                glidein_dict[(factory_pool_node, glidename, my_identity_at_factory_pool)] = factory_glidein_dict[
                    glidename
                ]

    # schedd
    condorq_format_list = elementDescript.merged_data["JobMatchAttrs"]
    condorq_dict = glideinFrontendLib.getCondorQ(
        elementDescript.merged_data["JobSchedds"], elementDescript.merged_data["JobQueryExpr"], condorq_format_list
    )

    # Match
    job_glideins_dict = {}
    glidein_jobs_dict = {}
    for glidename in glidein_dict:
        jobs = []
        sk = condorq_dict.keys()
        sk.sort()
        for schedd in sk:
            condorq = condorq_dict[schedd]
            condorq_data = condorq.fetchStored()
            jk = condorq_data.keys()
            jk.sort()
            for jid in jk:
                job = condorq_data[jid]
                t = "%s#%i.%i" % (schedd, jid[0], jid[1])
                if t not in job_glideins_dict:
                    job_glideins_dict[t] = []
                if eval(match_obj):
                    jobs.append(t)
                    job_glideins_dict[t].append(glidename[1])
                pass
            pass
        glidein_jobs_dict[glidename[1]] = jobs
        pass

    # Print
    print("Glideins")
    print("============================================================")
    gs = glidein_jobs_dict.keys()
    gs.sort()
    for glidename in gs:
        print("%-60s %s" % (glidename, glidein_jobs_dict[glidename]))

    print()
    print("Jobs")
    print("============================================================")
    js = job_glideins_dict.keys()
    js.sort()
    for job in js:
        print("%-40s %s" % (job, job_glideins_dict[job]))

    return


############################################################
def main(work_dir, group_name):
    elementDescript = glideinFrontendConfig.ElementMergedDescript(work_dir, group_name)

    paramsDescript = glideinFrontendConfig.ParamsDescript(work_dir, group_name)
    signatureDescript = glideinFrontendConfig.GroupSignatureDescript(work_dir, group_name)

    # set the condor configuration and GSI setup globally, so I don't need to worry about it later on
    os.environ["CONDOR_CONFIG"] = elementDescript.frontend_data["CondorConfig"]
    os.environ["_CONDOR_CERTIFICATE_MAPFILE"] = elementDescript.element_data["MapFile"]
    os.environ["X509_USER_PROXY"] = elementDescript.frontend_data["ClassAdProxy"]

    try:
        do_work(elementDescript, paramsDescript, signatureDescript)
    except KeyboardInterrupt:
        print("Received signal...exit")
    except Exception:
        tb = traceback.format_exception(sys.exc_info()[0], sys.exc_info()[1], sys.exc_info()[2])
        print("Exception occurred: %s" % tb)


############################################################
#
# S T A R T U P
#
############################################################

if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])
