#!/usr/bin/env python
# encoding: utf-8
"""
wmagent-mod-config.py

Created by Dave Evans on 2011-01-19.
Copyright (c) 2011 Fermilab. All rights reserved.

Note: this script is also used by ASO deployment, see:
https://github.com/dmwm/deployment/blob/master/asyncstageout/manage#L246
"""

import getopt
import importlib
import os
import socket
import sys
import traceback
from urllib.parse import urlparse

from WMCore.Configuration import saveConfigurationFile
from WMCore.Lexicon import splitCouchServiceURL

help_message = "The help message goes here."


class Usage(Exception):
    def __init__(self, msg):
        super(Usage, self).__init__(msg)
        self.msg = msg


### Code extracted from:
# https://stackoverflow.com/questions/5574702/how-to-print-to-stderr-in-python
def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def importConfigTemplate(filename):
    """
    _importConfigTemplate_

    Given filename, load it and grab the configuration object from it

    """

    cfgBaseName = os.path.basename(filename).replace(".py", "")
    cfgDirName = os.path.dirname(filename)
    if not cfgDirName:
        modSpecs = importlib.machinery.PathFinder().find_spec(cfgBaseName)
    else:
        modSpecs = importlib.machinery.PathFinder().find_spec(cfgBaseName, [cfgDirName])
    try:
        modRef = modSpecs.loader.load_module(cfgBaseName)
    except Exception as ex:
        msg = "Unable to load Configuration File: %s\n Due to error:\n" % filename
        msg += str(ex)
        msg += str(traceback.format_exc())
        raise RuntimeError(msg)
    config = getattr(modRef, 'config', None)
    if config is None:
        msg = "No config attribute found in %s" % filename
        raise RuntimeError(msg)
    return config


def saveConfiguration(configObject, outputPath):
    """
    _saveConfiguration_

    Save the configuration to the output path provided

    """
    saveConfigurationFile(configObject, outputPath)
    os.chmod(outputPath, 0o600)


def modifyConfiguration(config, **args):
    """
    _modifyConfiguration_

    Given the dictionary of key: value, look up the entry matching the key in the configuration
    and set it to that value in the config

    """
    mapping = {
        "coredb_url": [("CoreDatabase", "connectUrl")],
        "mysql_url": [("CoreDatabase", "connectUrl")],
        "couch_url": [("JobStateMachine", "couchurl"), ("ACDC", "couchurl"), ("WorkQueueManager", "couchurl")],
        "working_dir": [("General", "workDir")],
    }

    for k, v in args.items():
        parameters = mapping.get(k, [])
        for p in parameters:
            if hasattr(config, p[0]):
                section = getattr(config, p[0])
                setattr(section, p[1], v)
    for component in config.listComponents_():
        compCfg = getattr(config, component)
        compCfg.componentDir = "%s/%s" % (args['working_dir'], component)
        if component == "JobCreator":
            compCfg.jobCacheDir = "%s/%s/JobCache" % (args['working_dir'], component)
        elif component == "JobAccountant":
            compCfg.specDir = "%s/%s/SpecCache" % (args['working_dir'], component)

    localhost = socket.getfqdn()
    for webapp in config.listWebapps_():
        compCfg = getattr(config, webapp)
        compCfg.componentDir = "%s/%s" % (args['working_dir'], webapp)
        # Webtool REST service host location
        compCfg.Webtools.host = localhost

    # hostname locations
    config.Agent.hostName = localhost

    # configure MySQL specific settings
    if args.get('mysql_socket') and hasattr(config, "CoreDatabase"):
        config.CoreDatabase.socket = args['mysql_socket']

    if hasattr(config, "General"):
        config.General.central_logdb_url = args["central_logdb_url"]
        config.General.centralWMStatsURL = args["wmstats_url"]
        # t0 agent doesn't have reqmgr2_url
        if "reqmgr2_url" in args:
            config.General.ReqMgr2ServiceURL = args["reqmgr2_url"]

    if hasattr(config, "WorkloadSummary"):
        splited = args["workload_summary_url"].strip().rstrip('/').rsplit('/', 1)
        config.WorkloadSummary.couchurl = splited[0]
        config.WorkloadSummary.database = splited[1]

    if hasattr(config, "TaskArchiver"):
        splited = args["workload_summary_url"].strip().rstrip('/').rsplit('/', 1)
        config.TaskArchiver.workloadSummaryCouchURL = splited[0]
        config.TaskArchiver.workloadSummaryCouchDBName = splited[1]
        config.TaskArchiver.localCouchURL = "%s/%s" % (config.JobStateMachine.couchurl,
                                                       config.JobStateMachine.couchDBName)
        config.TaskArchiver.localQueueURL = "%s/%s" % (config.WorkQueueManager.couchurl,
                                                       config.WorkQueueManager.dbname)
        config.TaskArchiver.localWMStatsURL = "%s/%s" % (config.JobStateMachine.couchurl,
                                                         config.JobStateMachine.jobSummaryDBName)
        config.TaskArchiver.dqmUrl = args["dqm_url"]

    if hasattr(config, "ACDC"):
        # t0 agent may not have acdc_url
        if "acdc_url" in args:
            config.ACDC.couchurl, config.ACDC.database = splitCouchServiceURL(args["acdc_url"])

    if hasattr(config, "DBS3Upload"):
        config.DBS3Upload.dbsUrl = args["dbs3_url"]

    if hasattr(config, "AsyncTransfer"):
        config.CoreDatabase.connectUrl = '%s/asynctransfer_agent' % args["couch_url"]
        config.AsyncTransfer.couch_instance = args["couch_url"]
        config.AsyncTransfer.data_source = args["couch_url"]
        config.AsyncTransfer.serverDN = args["host_dn"]
        config.AsyncTransfer.couch_user_monitoring_instance = args["couch_url"]
        config.AsyncTransfer.opsProxy = args["ops_proxy"]
        config.AsyncTransfer.config_couch_instance = args["couch_url"]
        config.AsyncTransfer.cache_area = args["ufc_service_url"]

    if hasattr(config, "Analytics"):
        config.Analytics.couch_user_monitoring_instance = args["couch_url"]
        config.Analytics.amq_auth_file = args["amq_auth_file"]
        config.Analytics.couch_instance = args["couch_url"]
        config.Analytics.config_couch_instance = args["couch_url"]

    if hasattr(config, "FilesCleaner"):
        config.FilesCleaner.couch_instance = args["couch_url"]
        config.FilesCleaner.config_couch_instance = args["couch_url"]
        config.FilesCleaner.opsProxy = args["ops_proxy"]

    if hasattr(config, "Statistics"):
        config.Statistics.couch_statinstance = args["couch_url"]
        config.Statistics.couch_instance = args["couch_url"]
        config.Statistics.config_couch_instance = args["couch_url"]

    if hasattr(config, "DBSPublisher"):
        config.DBSPublisher.couch_instance = args["couch_url"]
        config.DBSPublisher.serverDN = args["host_dn"]
        config.DBSPublisher.cache_area = args["ufc_service_url"]
        config.DBSPublisher.opsProxy = args["ops_proxy"]
        config.DBSPublisher.config_couch_instance = args["couch_url"]

    if hasattr(config, "RetryManager"):
        # only needed for CRAB instance
        if "couch_url" in args:
            config.RetryManager.couch_instance = args["couch_url"]
        # only needed for CRAB instance
        if "ops_proxy" in args:
            config.RetryManager.opsProxy = args["ops_proxy"]

    if hasattr(config, "UserFileCache"):
        config.UserFileCache.Webtools.host = args["ufc_hostname"]
        config.UserFileCache.Webtools.port = int(args["ufc_port"])
        config.UserFileCache.views.active.userfilecache.serviceURL = args["ufc_service_url"]
        config.UserFileCache.userCacheDir = args["ufc_cachedir"]

    # custom test global workqueue
    if hasattr(config, "WorkQueueManager") and getattr(config.WorkQueueManager, "level", None) == 'GlobalQueue':
        urlObj = urlparse(config.WorkQueueManager.couchurl)
        if urlObj.port:
            couchPort = ":%s" % urlObj.port
        else:
            couchPort = ""
        config.WorkQueueManager.queueParams["QueueURL"] = "%s://%s%s" % (urlObj.scheme,
                                                                         config.Agent.hostName, couchPort)
        if args.get('workqueue_db_name'):
            config.WorkQueueManager.dbname = args['workqueue_db_name']
            config.WorkQueueManager.inboxDatabase = config.WorkQueueManager.dbname + '_inbox'
        config.WorkQueueManager.queueParams['WMStatsCouchUrl'] = "%s/%s" % (config.WorkQueueManager.couchurl.rstrip(),
                                                                            config.WorkQueueManager.wmstatDBName)
        config.WorkQueueManager.queueParams['ReqMgrServiceURL'] = args["reqmgr2_url"]
        config.WorkQueueManager.queueParams['RequestDBURL'] = args["requestcouch_url"]
        config.WorkQueueManager.queueParams['central_logdb_url'] = config.General.central_logdb_url
        config.WorkQueueManager.queueParams['log_reporter'] = "TEST_global_queue"

    # Custom Local WorkQueue
    if hasattr(config, "WorkQueueManager") and getattr(config.WorkQueueManager, "level", None) == 'LocalQueue':
        # set QueueURL param - this need to be unique between local WQs
        urlObj = urlparse(config.WorkQueueManager.couchurl)
        if urlObj.port:
            couchPort = ":%s" % urlObj.port
        else:
            couchPort = ""
        config.WorkQueueManager.queueParams["QueueURL"] = "%s://%s%s" % (urlObj.scheme,
                                                                         config.Agent.hostName, couchPort)

        if args.get("global_workqueue_url", None):
            config.WorkQueueManager.queueParams["ParentQueueCouchUrl"] = args["global_workqueue_url"]

        if args.get('workqueue_db_name'):
            config.WorkQueueManager.dbname = args['workqueue_db_name']
            config.WorkQueueManager.inboxDatabase = config.WorkQueueManager.dbname + '_inbox'

        config.WorkQueueManager.queueParams['RequestDBURL'] = args["requestcouch_url"]
        config.WorkQueueManager.queueParams['central_logdb_url'] = config.General.central_logdb_url
        config.WorkQueueManager.queueParams['log_reporter'] = config.Agent.hostName
        config.WorkQueueManager.rucioUrl = args["rucio_host"]
        config.WorkQueueManager.rucioAuthUrl = args["rucio_auth"]

    # custom AnalyticsDataCollector
    if hasattr(config, "AnalyticsDataCollector"):
        config.AnalyticsDataCollector.localCouchURL = "%s/%s" % (config.JobStateMachine.couchurl,
                                                                 config.JobStateMachine.couchDBName)
        config.AnalyticsDataCollector.localQueueURL = "%s/%s" % (config.WorkQueueManager.couchurl,
                                                                 config.WorkQueueManager.dbname)
        config.AnalyticsDataCollector.localWMStatsURL = "%s/%s" % (config.JobStateMachine.couchurl,
                                                                   config.JobStateMachine.jobSummaryDBName)
        config.AnalyticsDataCollector.centralRequestDBURL = args["requestcouch_url"]
        config.AnalyticsDataCollector.RequestCouchApp = "ReqMgr"

    # custom RucioInjector
    if hasattr(config, "RucioInjector"):
        config.RucioInjector.rucioAccount = args["rucio_account"]
        config.RucioInjector.rucioUrl = args["rucio_host"]
        config.RucioInjector.rucioAuthUrl = args["rucio_auth"]
        # define a different expression for container rule replication in testbed
        if "-int.cern.ch" in args["rucio_host"]:
            config.RucioInjector.containerDiskRuleRSEExpr = "(tier=2|tier=1)&cms_type=int&rse_type=DISK"

    # custom AgentStatusWatcher
    if hasattr(config, "AgentStatusWatcher"):
        config.AgentStatusWatcher.grafanaToken = "Bearer %s" % args["grafana_token"]
        if args.get("amq_credentials"):
            # e.g.: user@@@pass@@@topic
            amquser, amqpass, amqtopic = args["amq_credentials"].split('@@@')
            config.AgentStatusWatcher.userAMQ = amquser
            config.AgentStatusWatcher.passAMQ = amqpass
            config.AgentStatusWatcher.topicAMQ = amqtopic
            config.AgentStatusWatcher.enableAMQ = True

    # custom WMArchiveReporter
    if hasattr(config, "ArchiveDataReporter"):
        config.ArchiveDataReporter.WMArchiveURL = args["wmarchive_url"]

    # custom WMArchiveReporter
    if hasattr(config, "WorkflowUpdater"):
        config.WorkflowUpdater.dbsUrl = args.get("dbs3_reader_url", None)
        config.WorkflowUpdater.rucioUrl = args["rucio_host"]
        config.WorkflowUpdater.rucioAuthUrl = args["rucio_auth"]
        if args.get("wmstats_url", None):
            config.WorkflowUpdater.wmstatsUrl = args["wmstats_url"].replace("couchdb/wmstats", "wmstatsserver")
        if args.get("mspileup_url", None):
            config.WorkflowUpdater.msPileupUrl = args["mspileup_url"]


    return config


def main(argv=None):
    if argv is None:
        argv = sys.argv

    inputFile = None
    outputFile = None
    parameters = {}

    try:
        try:
            ### For the record, options used by ASO deployment are:
            # input, output, couch_url, working_dir, host_dn, wmstats_url
            # ufc_service_url, ops_proxy and amq_auth_file
            ### in case we ever decide to clean this up
            opts, _args = getopt.getopt(argv[1:], "h",
                                        ["help", "input=", "output=", "mysql_socket=", "mysql_url=", "couch_url=",
                                         "working_dir=", "global_workqueue_url=",
                                         "workqueue_db_name=", "cs_hostname=", "cs_port=", "host_dn=",
                                         "sb_hostname=", "sb_port=", "sb_basepath=", "ufc_hostname=", "ufc_port=",
                                         "ufc_service_url=", "ufc_cachedir=", "client_mapping=",
                                         "workload_summary_url=", "coredb_url=", "wmstats_url=", "ops_proxy=",
                                         "reqmgr2_url=", "acdc_url=", "amq_auth_file=", "dbs3_url=", "dbs3_reader_url=",
                                         "dqm_url=", "grafana_token=", "requestcouch_url=", "central_logdb_url=",
                                         "wmarchive_url=", "amq_credentials=",
                                         "rucio_account=", "rucio_host=", "rucio_auth=", "mspileup_url="])

        except getopt.error as msg:
            raise Usage(msg)

        # option processing
        for option, value in opts:
            if option in ("-h", "--help"):
                raise Usage(help_message)
            if option == "--output":
                outputFile = value
            if option == "--input":
                inputFile = value
            if option in ('--mysql_socket', '--mysql_url', '--coredb_url', '--couch_url',
                          '--working_dir',
                          '--global_workqueue_url', '--workqueue_db_name',
                          '--cs_hostname', '--cs_port', '--host_dn',
                          '--sb_hostname', '--sb_port', '--sb_basepath', '--ufc_hostname',
                          '--ufc_port', '--ufc_service_url', '--ufc_cachedir',
                          '--client_mapping', '--workload_summary_url',
                          '--wmstats_url', '--ops_proxy', '--reqmgr2_url', '--acdc_url',
                          '--amq_auth_file', '--dbs3_url', '--dbs3_reader_url', '--dqm_url',
                          '--grafana_token', '--requestcouch_url', '--central_logdb_url',
                          '--wmarchive_url', '--amq_credentials',
                          '--rucio_account', '--rucio_host', '--rucio_auth', '--mspileup_url'):
                parameters[option[2:]] = value


    except Usage as err:
        eprint(sys.argv[0].split("/")[-1] + ": " + str(err.msg))
        eprint("\t for help use --help")
        return 2

    try:
        cfg = importConfigTemplate(inputFile)
    except Exception as ex:
        msg = "Failed to import template config: %s\n" % inputFile
        msg += str(ex)
        eprint(msg)
        return 3
    try:
        cfg = modifyConfiguration(cfg, **parameters)
    except Exception as ex:
        msg = "Error modifying configuration:\n %s" % str(ex)
        eprint(msg)
        eprint(traceback.format_exc())
        return 4
    try:
        saveConfiguration(cfg, outputFile)
    except Exception as ex:
        msg = "Error saving output configuration file:\n %s\n" % outputFile
        msg += str(ex)
        eprint(msg)
        return 5


if __name__ == "__main__":
    sys.exit(main())
