#!/usr/bin/env python

import os
import sys
import json
from pprint import pformat

from WMCore.Agent.DefaultConfig import DEFAULT_AGENT_CONFIG
from WMCore.Configuration import loadConfigurationFile
from WMCore.Services.ReqMgrAux.ReqMgrAux import ReqMgrAux

if __name__ == "__main__":
    if len(sys.argv) > 2:
        print("Too many arguments provided: %s" % sys.argv)
        print("Usage: python wmagent-upload-config [optional override]")
        print("Example: python wmagent-upload-config '{\"MaxRetries\":2}'")
        sys.exit(0)
    elif len(sys.argv) == 2:
        override = json.loads(sys.argv[1])
    else:
        override = {}

    wmConfig = loadConfigurationFile(os.environ['WMAGENT_CONFIG'])
    reqMgrAux = ReqMgrAux(wmConfig.General.ReqMgr2ServiceURL)

    DEFAULT_AGENT_CONFIG.update(override)
    print("Pushing the following agent configuration:\n%s" % pformat(DEFAULT_AGENT_CONFIG))
    if not reqMgrAux.updateWMAgentConfig(wmConfig.Agent.hostName, DEFAULT_AGENT_CONFIG):
        # then the document does not exist, create it
        res = reqMgrAux.postWMAgentConfig(wmConfig.Agent.hostName, DEFAULT_AGENT_CONFIG)
        print("Agent config successfully created: %s" % res.get("ok", False))
