""" Class that contains client access to the StorageManagerDB handler.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
__RCSID__ = "$Id$"

import six
import random
import errno

from DIRAC import S_OK, S_ERROR, gLogger
from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getVOForGroup
from DIRAC.Core.Base.Client import Client, createClient
from DIRAC.Core.Utilities.DErrno import cmpError
from DIRAC.Core.Utilities.Proxy import UserProxy
from DIRAC.DataManagementSystem.Client.DataManager import DataManager
from DIRAC.DataManagementSystem.Utilities.DMSHelpers import DMSHelpers
from DIRAC.Resources.Storage.StorageElement import StorageElement


def getFilesToStage(lfnList, jobState=None, checkOnlyTapeSEs=None, jobLog=None):
  """ Utility that returns out of a list of LFNs those files that are offline,
      and those for which at least one copy is online
  """
  if not lfnList:
    return S_OK({'onlineLFNs': [], 'offlineLFNs': {}, 'failedLFNs': [], 'absentLFNs': {}})

  dm = DataManager()
  if isinstance(lfnList, six.string_types):
    lfnList = [lfnList]

  lfnListReplicas = dm.getReplicasForJobs(lfnList, getUrl=False)
  if not lfnListReplicas['OK']:
    return lfnListReplicas

  offlineLFNsDict = {}
  onlineLFNs = {}
  offlineLFNs = {}
  absentLFNs = {}
  failedLFNs = set()
  if lfnListReplicas['Value']['Failed']:
    # Check if files are not existing
    for lfn, reason in lfnListReplicas['Value']['Failed'].items():
      # FIXME: awful check until FC returns a proper error
      if cmpError(reason, errno.ENOENT) or 'No such file' in reason:
        # The file doesn't exist, job must be Failed
        # FIXME: it is not possible to return here an S_ERROR(), return the message only
        absentLFNs[lfn] = S_ERROR(errno.ENOENT, 'File not in FC')['Message']
    if absentLFNs:
      return S_OK({'onlineLFNs': list(onlineLFNs),
                   'offlineLFNs': offlineLFNsDict,
                   'failedLFNs': list(failedLFNs),
                   'absentLFNs': absentLFNs})
    return S_ERROR("Failures in getting replicas")

  lfnListReplicas = lfnListReplicas['Value']['Successful']
  # If a file is reported here at a tape SE, it is not at a disk SE as we use disk in priority
  # We shall check all file anyway in order to make sure they exist
  seToLFNs = dict()
  for lfn, ses in lfnListReplicas.items():
    for se in ses:
      seToLFNs.setdefault(se, list()).append(lfn)

  if seToLFNs:
    if jobState:
      # Get user name and group from the job state
      userName = jobState.getAttribute('Owner')
      if not userName['OK']:
        return userName
      userName = userName['Value']

      userGroup = jobState.getAttribute('OwnerGroup')
      if not userGroup['OK']:
        return userGroup
      userGroup = userGroup['Value']
    else:
      userName = None
      userGroup = None
    # Check whether files are Online or Offline, or missing at SE
    result = _checkFilesToStage(seToLFNs, onlineLFNs, offlineLFNs, absentLFNs,  # pylint: disable=unexpected-keyword-arg
                                checkOnlyTapeSEs=checkOnlyTapeSEs, jobLog=jobLog,
                                proxyUserName=userName,
                                proxyUserGroup=userGroup,
                                executionLock=True)

    if not result['OK']:
      return result
    failedLFNs = set(lfnList) - set(onlineLFNs) - set(offlineLFNs) - set(absentLFNs)

    # Get the online SEs
    dmsHelper = DMSHelpers()
    onlineSEs = set(se for ses in onlineLFNs.values() for se in ses)
    onlineSites = set(dmsHelper.getLocalSiteForSE(se).get('Value') for se in onlineSEs) - {None}
    for lfn in offlineLFNs:
      ses = offlineLFNs[lfn]
      if len(ses) == 1:
        # No choice, let's go
        offlineLFNsDict.setdefault(ses[0], list()).append(lfn)
        continue
      # Try and get an SE at a site already with online files
      found = False
      if onlineSites:
        # If there is at least one online site, select one
        for se in ses:
          site = dmsHelper.getLocalSiteForSE(se)
          if site['OK']:
            if site['Value'] in onlineSites:
              offlineLFNsDict.setdefault(se, list()).append(lfn)
              found = True
              break
      # No online site found in common, select randomly
      if not found:
        offlineLFNsDict.setdefault(random.choice(ses), list()).append(lfn)

  return S_OK({'onlineLFNs': list(onlineLFNs),
               'offlineLFNs': offlineLFNsDict,
               'failedLFNs': list(failedLFNs),
               'absentLFNs': absentLFNs,
               'onlineSites': onlineSites})


def _checkFilesToStage(seToLFNs, onlineLFNs, offlineLFNs, absentLFNs,
                       checkOnlyTapeSEs=None, jobLog=None,
                       proxyUserName=None,
                       proxyUserGroup=None,
                       executionLock=None):
  """
  Checks on SEs whether the file is NEARLINE or ONLINE
  onlineLFNs, offlineLFNs and absentLFNs are modified to contain the files found online
  If checkOnlyTapeSEs is True, disk replicas are not checked
  As soon as a replica is found Online for a file, no further check is made
  """
  # Only check on storage if it is a tape SE
  if jobLog is None:
    logger = gLogger
  else:
    logger = jobLog
  if checkOnlyTapeSEs is None:
    # Default value is True
    checkOnlyTapeSEs = True

  failed = {}
  for se, lfnsInSEList in seToLFNs.items():
    # If we have found already all files online at another SE, no need to check the others
    # but still we want to set the SE as Online if not a TapeSE
    vo = getVOForGroup(proxyUserGroup)
    seObj = StorageElement(se, vo=vo)
    status = seObj.getStatus()
    if not status['OK']:
      return status
    tapeSE = status['Value']['TapeSE']
    diskSE = status['Value']['DiskSE']
    # If requested to check only Tape SEs and the file is at a diskSE, we guess it is Online...
    filesToCheck = []
    for lfn in lfnsInSEList:
      # If the file had already been found accessible at an SE, only check that this one is on disk
      diskIsOK = checkOnlyTapeSEs or (lfn in onlineLFNs)
      if diskIsOK and diskSE:
        onlineLFNs.setdefault(lfn, []).append(se)
      elif not diskIsOK or (tapeSE and (lfn not in onlineLFNs)):
        filesToCheck.append(lfn)
    if not filesToCheck:
      continue

    # We have to use a new SE object because it caches the proxy!
    with UserProxy(proxyUserName=proxyUserName,
                   proxyUserGroup=proxyUserGroup,
                   executionLock=executionLock) as proxyResult:
      if proxyResult['OK']:
        fileMetadata = StorageElement(se, vo=vo).getFileMetadata(filesToCheck)
      else:
        fileMetadata = proxyResult

    if not fileMetadata['OK']:
      failed[se] = dict.fromkeys(filesToCheck, fileMetadata['Message'])
    else:
      if fileMetadata['Value']['Failed']:
        failed[se] = fileMetadata['Value']['Failed']
      # is there at least one replica online?
      for lfn, mDict in fileMetadata['Value']['Successful'].items():
        # SRM returns Cached, but others may only return Accessible
        if mDict.get('Cached', mDict['Accessible']):
          onlineLFNs.setdefault(lfn, []).append(se)
        elif tapeSE:
          # A file can be staged only at Tape SE
          offlineLFNs.setdefault(lfn, []).append(se)
        else:
          # File not available at a diskSE... we shall retry later
          pass

  # Doesn't matter if some files are Offline if they are also online
  for lfn in set(offlineLFNs) & set(onlineLFNs):
    offlineLFNs.pop(lfn)

  # If the file was found staged, ignore possible errors, but print out errors
  for se, failedLfns in list(failed.items()):
    logger.error("Errors when getting files metadata", 'at %s' % se)
    for lfn, reason in list(failedLfns.items()):
      if lfn in onlineLFNs:
        logger.warn(reason, 'for %s, but there is an online replica' % lfn)
        failed[se].pop(lfn)
      else:
        logger.error(reason, 'for %s, no online replicas' % lfn)
        if cmpError(reason, errno.ENOENT):
          absentLFNs.setdefault(lfn, []).append(se)
          failed[se].pop(lfn)
    if not failed[se]:
      failed.pop(se)
  # Find the files that do not exist at SE
  if failed:
    logger.error("Error getting metadata", "for %d files" %
                 len(set(lfn for lfnList in failed.values() for lfn in lfnList)))

  for lfn in absentLFNs:
    seList = absentLFNs[lfn]
    # FIXME: it is not possible to return here an S_ERROR(), return the message only
    absentLFNs[lfn] = S_ERROR(errno.ENOENT, "File not at %s" % ','.join(sorted(seList)))['Message']
  # Format the error for absent files
  return S_OK()


@createClient('StorageManagement/StorageManager')
class StorageManagerClient(Client):
  """ This is the client to the StorageManager service, so even if it is not seen, it exposes all its RPC calls
  """

  def __init__(self, **kwargs):
    super(StorageManagerClient, self).__init__(**kwargs)
    self.setServer('StorageManagement/StorageManager')
