"""Module collecting functions dealing with the GLUE2 information schema

:author: A.Sailer

Known problems:

 * ARC CEs do not seem to publish wall or CPU time per queue anywhere
 * There is no consistency between which memory information is provided where,
   execution environment vs. information for a share
 * Some execution environment IDs are used more than once

Print outs with "SCHEMA PROBLEM" point -- in my opinion -- to errors in the
published information, like a foreign key pointing to non-existent entry.

"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import six
from pprint import pformat

from DIRAC import gLogger, gConfig
from DIRAC.ConfigurationSystem.Client.Helpers.Resources import getCESiteMapping, getGOCSiteName
from DIRAC.Core.Utilities.ReturnValues import S_OK, S_ERROR
from DIRAC.Core.Utilities.List import breakListIntoChunks

__RCSID__ = "$Id$"


sLog = gLogger.getSubLogger(__name__)


def __ldapsearchBDII(*args, **kwargs):
  """ wrap `DIRAC.Core.Utilities.Grid.ldapsearchBDII` to avoid circular import """
  from DIRAC.Core.Utilities.Grid import ldapsearchBDII
  return ldapsearchBDII(*args, **kwargs)


def getGlue2CEInfo(vo, host):
  """ call ldap for GLUE2 and get information

  :param str vo: Virtual Organisation
  :param str host: host to query for information
  :returns: result structure with result['Value'][siteID]['CEs'][ceID]['Queues'][queueName]. For
               each siteID, ceID, queueName all the GLUE2 parameters are retrieved
  """

  # get all Policies allowing given VO
  filt = "(&(objectClass=GLUE2Policy)(|(GLUE2PolicyRule=VO:%s)(GLUE2PolicyRule=vo:%s)))" % (vo, vo)
  polRes = __ldapsearchBDII(filt=filt, attr=None, host=host, base="o=glue", selectionString="GLUE2")

  if not polRes['OK']:
    return S_ERROR("Failed to get policies for this VO")
  polRes = polRes['Value']

  sLog.notice("Found %s policies for this VO %s" % (len(polRes), vo))
  # get all shares for this policy
  # create an or'ed list of all the shares and then call the search
  listOfSitesWithPolicies = set()
  shareFilter = ''
  for policyValues in polRes:
    # skip entries without GLUE2DomainID in the DN because we cannot associate them to a site
    if 'GLUE2DomainID' not in policyValues['attr']['dn']:
      continue
    shareID = policyValues['attr'].get('GLUE2MappingPolicyShareForeignKey', None)
    policyID = policyValues['attr']['GLUE2PolicyID']
    siteName = policyValues['attr']['dn'].split('GLUE2DomainID=')[1].split(',', 1)[0]
    listOfSitesWithPolicies.add(siteName)
    if shareID is None:  # policy not pointing to ComputingInformation
      sLog.debug("Policy %s does not point to computing information" % (policyID,))
      continue
    sLog.verbose("%s policy %s pointing to %s " % (siteName, policyID, shareID))
    sLog.debug("Policy values:\n%s" % pformat(policyValues))
    shareFilter += '(GLUE2ShareID=%s)' % shareID

  filt = '(&(objectClass=GLUE2Share)(|%s))' % shareFilter
  shareRes = __ldapsearchBDII(filt=filt, attr=None, host=host, base="o=glue", selectionString="GLUE2")
  if not shareRes['OK']:
    sLog.error("Could not get share information", shareRes['Message'])
    return shareRes
  shareInfoLists = {}
  for shareInfo in shareRes['Value']:
    if 'GLUE2DomainID' not in shareInfo['attr']['dn']:
      continue
    if 'GLUE2ComputingShare' not in shareInfo['objectClass']:
      sLog.debug('Share %r is not a ComputingShare: \n%s' % (shareID, pformat(shareInfo)))
      continue
    sLog.debug("Found computing share:\n%s" % pformat(shareInfo))
    siteName = shareInfo['attr']['dn'].split('GLUE2DomainID=')[1].split(',', 1)[0]
    shareInfoLists.setdefault(siteName, []).append(shareInfo['attr'])

  siteInfo = __getGlue2ShareInfo(host, shareInfoLists)
  if not siteInfo['OK']:
    sLog.error("Could not get CE info for", "%s: %s" % (shareID, siteInfo['Message']))
    return siteInfo
  siteDict = siteInfo['Value']
  sLog.debug("Found Sites:\n%s" % pformat(siteDict))
  sitesWithoutShares = set(siteDict) - listOfSitesWithPolicies
  if sitesWithoutShares:
    sLog.error("Found some sites without any shares", pformat(sitesWithoutShares))
  else:
    sLog.notice("Found information for all known sites")

  # remap siteDict to assign CEs to known sites, in case their names differ from the "gocdb name" in
  # the CS.
  newSiteDict = {}
  ceSiteMapping = getCESiteMapping().get('Value', {})
  # FIXME: pylint thinks siteDict is a tuple, so we cast
  for siteName, infoDict in dict(siteDict).items():
    for ce, ceInfo in infoDict.get('CEs', {}).items():
      ceSiteName = ceSiteMapping.get(ce, siteName)
      gocSiteName = getGOCSiteName(ceSiteName).get('Value', siteName)
      newSiteDict.setdefault(gocSiteName, {}).setdefault('CEs', {})[ce] = ceInfo

  return S_OK(newSiteDict)


def __getGlue2ShareInfo(host, shareInfoLists):
  """ get information from endpoints, which are the CE at a Site

  :param str host: BDII host to query
  :param dict shareInfoDict: dictionary of GLUE2 parameters belonging to the ComputingShare
  :returns: result structure S_OK/S_ERROR
  """
  executionEnvironments = []
  for _siteName, shareInfoDicts in shareInfoLists.items():
    for shareInfoDict in shareInfoDicts:
      executionEnvironment = shareInfoDict.get('GLUE2ComputingShareExecutionEnvironmentForeignKey', [])
      if not executionEnvironment:
        sLog.error('No entry for GLUE2ComputingShareExecutionEnvironmentForeignKey', pformat(shareInfoDict))
        continue
      if isinstance(executionEnvironment, six.string_types):
        executionEnvironment = [executionEnvironment]
      executionEnvironments.extend(executionEnvironment)
  resExeInfo = __getGlue2ExecutionEnvironmentInfo(host, executionEnvironments)
  if not resExeInfo['OK']:
    sLog.error("Cannot get execution environment info for:",
               str(executionEnvironments)[:100] + "  " + resExeInfo['Message'])
    return resExeInfo
  exeInfos = resExeInfo['Value']

  siteDict = {}
  for siteName, shareInfoDicts in shareInfoLists.items():
    siteDict[siteName] = {'CEs': {}}
    cesDict = siteDict[siteName]['CEs']
    for shareInfoDict in shareInfoDicts:
      ceInfo = {}
      ceInfo['MaxWaitingJobs'] = shareInfoDict.get('GLUE2ComputingShareMaxWaitingJobs', '-1')  # This is not used
      ceInfo['Queues'] = {}
      queueInfo = {}
      queueInfo['GlueCEStateStatus'] = shareInfoDict['GLUE2ComputingShareServingState']
      queueInfo['GlueCEPolicyMaxCPUTime'] = str(int(int(
          shareInfoDict.get('GLUE2ComputingShareMaxCPUTime', 86400)) / 60))
      queueInfo['GlueCEPolicyMaxWallClockTime'] = str(int(int(
          shareInfoDict.get('GLUE2ComputingShareMaxWallTime', 86400)) / 60))
      queueInfo['GlueCEInfoTotalCPUs'] = shareInfoDict.get('GLUE2ComputingShareMaxRunningJobs', '10000')
      queueInfo['GlueCECapability'] = ['CPUScalingReferenceSI00=2552']

      try:
        maxNOPfromCS = gConfig.getValue('/Resources/Computing/CEDefaults/GLUE2ComputingShareMaxSlotsPerJob_limit', 8)
        maxNOPfromGLUE = int(shareInfoDict.get('GLUE2ComputingShareMaxSlotsPerJob', 1))
        numberOfProcs = min(maxNOPfromGLUE, maxNOPfromCS)
        queueInfo['NumberOfProcessors'] = numberOfProcs
        if numberOfProcs != maxNOPfromGLUE:
          sLog.info('Limited NumberOfProcessors for', '%s from %s to %s' %
                    (siteName, maxNOPfromGLUE, numberOfProcs))
      except ValueError:
        sLog.error("Bad content for GLUE2ComputingShareMaxSlotsPerJob:",
                   siteName + ' ' + shareInfoDict.get('GLUE2ComputingShareMaxSlotsPerJob'))
        queueInfo['NumberOfProcessors'] = 1

      executionEnvironment = shareInfoDict.get('GLUE2ComputingShareExecutionEnvironmentForeignKey', [])
      if isinstance(executionEnvironment, six.string_types):
        executionEnvironment = [executionEnvironment]
      resExeInfo = __getGlue2ExecutionEnvironmentInfoForSite(siteName, executionEnvironment, exeInfos)
      if not resExeInfo['OK']:
        continue

      exeInfo = resExeInfo.get('Value')
      if not exeInfo:
        sLog.error('Using dummy values. Did not find information for execution environment', siteName)
        exeInfo = {'GlueHostMainMemoryRAMSize': '1999',  # intentionally identifiably dummy value
                   'GlueHostOperatingSystemVersion': '',
                   'GlueHostOperatingSystemName': '',
                   'GlueHostOperatingSystemRelease': '',
                   'GlueHostArchitecturePlatformType': 'x86_64',
                   'GlueHostBenchmarkSI00': '2500',  # needed for the queue to be used by the sitedirector
                   'MANAGER': 'manager:unknownBatchSystem',  # need some value for ARC
                   }
      else:
        sLog.info('Found information for execution environment for', siteName)

      # sometimes the time is still in hours
      maxCPUTime = int(queueInfo['GlueCEPolicyMaxCPUTime'])
      if maxCPUTime in [12, 24, 36, 48, 168]:
        queueInfo['GlueCEPolicyMaxCPUTime'] = str(maxCPUTime * 60)
        queueInfo['GlueCEPolicyMaxWallClockTime'] = str(int(queueInfo['GlueCEPolicyMaxWallClockTime']) * 60)

      ceInfo.update(exeInfo)
      shareEndpoints = shareInfoDict.get('GLUE2ShareEndpointForeignKey', [])
      if isinstance(shareEndpoints, six.string_types):
        shareEndpoints = [shareEndpoints]
      for endpoint in shareEndpoints:
        ceType = endpoint.rsplit('.', 1)[1]
        # get queue Name, in CREAM this is behind GLUE2entityOtherInfo...
        if ceType == 'CREAM':
          for otherInfo in shareInfoDict['GLUE2EntityOtherInfo']:
            if otherInfo.startswith('CREAMCEId'):
              queueName = otherInfo.split('/', 1)[1]
              # creamCEs are EOL soon, ignore any info they have
              if queueInfo.pop('NumberOfProcessors', 1) != 1:
                sLog.verbose('Ignoring MaxSlotsPerJob option for CreamCE', endpoint)

        # HTCondorCE, htcondorce
        elif ceType.lower().endswith('htcondorce'):
          ceType = 'HTCondorCE'
          queueName = 'condor'

        else:
          sLog.error('Unknown CE Type, please check the available information', ceType)
          continue

        queueInfo['GlueCEImplementationName'] = ceType
        ceName = endpoint.split('_', 1)[0]
        cesDict.setdefault(ceName, {})
        existingQueues = dict(cesDict[ceName].get('Queues', {}))
        existingQueues[queueName] = queueInfo
        ceInfo['Queues'] = existingQueues
        cesDict[ceName].update(ceInfo)

      # ARC CEs do not have endpoints, we have to try something else to get the information about the queue etc.
      try:
        if not shareEndpoints and shareInfoDict['GLUE2ShareID'].startswith('urn:ogf'):
          exeInfo = dict(exeInfo)  # silence pylint about tuples
          queueInfo['GlueCEImplementationName'] = 'ARC'
          managerName = exeInfo.pop('MANAGER', '').split(' ', 1)[0].rsplit(':', 1)[1]
          managerName = managerName.capitalize() if managerName == 'condor' else managerName
          queueName = 'nordugrid-%s-%s' % (managerName, shareInfoDict['GLUE2ComputingShareMappingQueue'])
          ceName = shareInfoDict['GLUE2ShareID'].split('ComputingShare:')[1].split(':')[0]
          cesDict.setdefault(ceName, {})
          existingQueues = dict(cesDict[ceName].get('Queues', {}))
          existingQueues[queueName] = queueInfo
          ceInfo['Queues'] = existingQueues
          cesDict[ceName].update(ceInfo)
      except Exception:
        sLog.error('Exception in ARC part for site:', siteName)

  return S_OK(siteDict)


def __getGlue2ExecutionEnvironmentInfo(host, executionEnvironments):
  """Find all the executionEnvironments.

  :param str host: BDII host to query
  :param list executionEnvironments: list of the execution environments to get some information from
  :returns: result of the ldapsearch for all executionEnvironments, Glue2 schema
  """
  listOfValues = []
  # break up to avoid argument list too long, it started failing at about 1900 entries
  for exeEnvs in breakListIntoChunks(executionEnvironments, 1000):
    exeFilter = ''
    for execEnv in exeEnvs:
      exeFilter += '(GLUE2ResourceID=%s)' % execEnv
    filt = "(&(objectClass=GLUE2ExecutionEnvironment)(|%s))" % exeFilter
    response = __ldapsearchBDII(filt=filt, attr=None, host=host, base="o=glue", selectionString="GLUE2")
    if not response['OK']:
      return response
    if not response['Value']:
      sLog.error("No information found for %s" % executionEnvironments)
      continue
    listOfValues += response['Value']
  if not listOfValues:
    return S_ERROR("No information found for executionEnvironments")
  return S_OK(listOfValues)


def __getGlue2ExecutionEnvironmentInfoForSite(sitename, foreignKeys, exeInfos):
  """Get the information about the execution environment for a specific site or ce or something.

  :param str sitename: Name of the site we are looking at
  :param list foreignKeys: list of ExecutionEnvironmentForeignkeys linked by the site
  :param list exeInfos: bdii list of dictionaries containing all the ExecutionEnvironment information for all sites
  :return: Dictionary with the information as required by the Bdii2CSagent for this site
  """
  # filter those that we want
  exeInfos = [exeInfo for exeInfo in exeInfos if exeInfo['attr']['GLUE2ResourceID'] in foreignKeys]
  # take the CE with the lowest MainMemory
  exeInfo = sorted(exeInfos, key=lambda k: int(k['attr']['GLUE2ExecutionEnvironmentMainMemorySize']))
  if not exeInfo:
    sLog.error('SCHEMA PROBLEM: Did not find execution info for site', sitename + ' and keys: ' + ' '.join(foreignKeys))
    return S_OK()
  sLog.debug("Found ExecutionEnvironments", pformat(exeInfo[0]))
  exeInfo = exeInfo[0]['attr']  # pylint: disable=unsubscriptable-object
  maxRam = exeInfo.get('GLUE2ExecutionEnvironmentMainMemorySize', '')
  architecture = exeInfo.get('GLUE2ExecutionEnvironmentPlatform', '')
  architecture = 'x86_64' if architecture == 'amd64' else architecture
  architecture = 'x86_64' if architecture == 'UNDEFINEDVALUE' else architecture
  architecture = 'x86_64' if "Intel(R) Xeon(R)" in architecture else architecture
  osFamily = exeInfo.get('GLUE2ExecutionEnvironmentOSFamily', '')  # e.g. linux
  osName = exeInfo.get('GLUE2ExecutionEnvironmentOSName', '')
  osVersion = exeInfo.get('GLUE2ExecutionEnvironmentOSVersion', '')
  manager = exeInfo.get('GLUE2ExecutionEnvironmentComputingManagerForeignKey', 'manager:unknownBatchSystem')
  # translate to Glue1 like keys, because that is used later on
  infoDict = {'GlueHostMainMemoryRAMSize': maxRam,
              'GlueHostOperatingSystemVersion': osName,
              'GlueHostOperatingSystemName': osFamily,
              'GlueHostOperatingSystemRelease': osVersion,
              'GlueHostArchitecturePlatformType': architecture.lower(),
              'GlueHostBenchmarkSI00': '2500',  # needed for the queue to be used by the sitedirector
              'MANAGER': manager,  # to create the ARC QueueName mostly
              }

  return S_OK(infoDict)
