""" Some utilities for FTS3...
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import random
import threading

from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations as opHelper
from DIRAC.DataManagementSystem.Client.DataManager import DataManager
from DIRAC.FrameworkSystem.Client.Logger import gLogger
from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
from DIRAC.Core.Utilities.ReturnValues import S_OK, S_ERROR
from DIRAC.ResourceStatusSystem.Client.ResourceStatus import ResourceStatus


def _checkSourceReplicas(ftsFiles, preferDisk=True):
  """ Check the active replicas

      :params ftsFiles: list of FT3Files
      :param preferDisk: (default True) prefer disk replicas if available
                          (see :py:meth:`DIRAC.DataManagementSystem.Client.DataManager.DataManager.getActiveReplicas`)

      :returns: Successful/Failed {lfn : { SE1 : PFN1, SE2 : PFN2 } , ... }
  """

  lfns = list(set([f.lfn for f in ftsFiles]))
  res = DataManager().getActiveReplicas(lfns, getUrl=False, preferDisk=preferDisk)

  return res


def selectUniqueSource(ftsFiles, fts3Plugin, allowedSources=None):
  """
      For a list of FTS3files object, select a random source, and group the files by source.

      We also return the FTS3Files for which we had problems getting replicas

      :param ftsFiles: list of FTS3File object
      :param fts3Plugin: plugin instance to use to chose between sources
      :param allowedSources: list of allowed sources


      :return:  S_OK(({ sourceSE: [ FTS3Files] }, {FTS3File: errors}))

  """

  _log = gLogger.getSubLogger("selectUniqueSource")

  # destGroup will contain for each target SE a dict { source : [list of FTS3Files] }
  groupBySource = {}

  # For all files, check which possible sources they have
  # If we specify allowedSources, don't restrict the choice to disk replicas
  res = _checkSourceReplicas(ftsFiles, preferDisk=not allowedSources)
  if not res['OK']:
    return res

  filteredReplicas = res['Value']

  # LFNs for which we failed to get replicas
  failedFiles = {}

  for ftsFile in ftsFiles:

    # If we failed to get the replicas, add the FTS3File to
    # the dictionnary
    if ftsFile.lfn in filteredReplicas['Failed']:
      errMsg = filteredReplicas['Failed'][ftsFile.lfn]
      failedFiles[ftsFile] = errMsg
      _log.debug("Failed to get active replicas", "%s,%s" %
                 (ftsFile.lfn, errMsg))
      continue

    replicaDict = filteredReplicas['Successful'][ftsFile.lfn]

    try:
      uniqueSource = fts3Plugin.selectSourceSE(ftsFile, replicaDict, allowedSources)
      groupBySource.setdefault(uniqueSource, []).append(ftsFile)
    except ValueError as e:
      _log.info('No allowed replica source for file', "%s: %s" % (ftsFile.lfn, repr(e)))
      continue

  return S_OK((groupBySource, failedFiles))


def groupFilesByTarget(ftsFiles):
  """
        For a list of FTS3files object, group the Files by target

        :param ftsFiles: list of FTS3File object
        :return: {targetSE : [ ftsFiles] } }

    """

  # destGroup will contain for each target SE a dict { possible source : transfer metadata }
  destGroup = {}

  for ftsFile in ftsFiles:
    destGroup.setdefault(ftsFile.targetSE, []).append(ftsFile)

  return S_OK(destGroup)


def getFTS3Plugin(vo=None):
  """
    Return an instance of the FTS3Plugin configured in the CS

    :param vo: vo config to look for
  """
  pluginName = opHelper(vo=vo).getValue('DataManagement/FTSPlacement/FTS3/FTS3Plugin', 'Default')

  objLoader = ObjectLoader()
  _class = objLoader.loadObject(
      'DataManagementSystem.private.FTS3Plugins.%sFTS3Plugin' %
      pluginName, '%sFTS3Plugin' % pluginName)

  if not _class['OK']:
    raise Exception(_class['Message'])

  fts3Plugin = _class['Value'](vo=vo)
  return fts3Plugin


threadLocal = threading.local()


class FTS3ServerPolicy(object):
  """
  This class manages the policy for choosing a server
  """

  def __init__(self, serverDict, serverPolicy="Random"):
    """
        Call the init of the parent, and initialize the list of FTS3 servers
    """

    self.log = gLogger.getSubLogger("FTS3ServerPolicy")

    self._serverDict = serverDict
    self._serverList = list(serverDict)
    self._maxAttempts = len(self._serverList)
    self._nextServerID = 0
    self._resourceStatus = ResourceStatus()

    methName = "_%sServerPolicy" % serverPolicy.lower()
    if not hasattr(self, methName):
      self.log.error('Unknown server policy %s. Using Random instead' % serverPolicy)
      methName = "_randomServerPolicy"

    self._policyMethod = getattr(self, methName)

  def _failoverServerPolicy(self, _attempt):
    """
       Returns always the server at a given position (normally the first one)

       :param attempt: position of the server in the list
    """
    if _attempt >= len(self._serverList):
      raise Exception(
          "FTS3ServerPolicy.__failoverServerPolicy: attempt to reach non existing server index")
    return self._serverList[_attempt]

  def _sequenceServerPolicy(self, _attempt):
    """
       Every time the this policy is called, return the next server on the list
    """

    fts3server = self._serverList[self._nextServerID]
    self._nextServerID = (self._nextServerID + 1) % len(self._serverList)
    return fts3server

  def _randomServerPolicy(self, _attempt):
    """
      return a server from shuffledServerList
    """

    if getattr(threadLocal, 'shuffledServerList', None) is None:
      threadLocal.shuffledServerList = self._serverList[:]
      random.shuffle(threadLocal.shuffledServerList)

    fts3Server = threadLocal.shuffledServerList[_attempt]

    if _attempt == self._maxAttempts - 1:
      random.shuffle(threadLocal.shuffledServerList)

    return fts3Server

  def _getFTSServerStatus(self, ftsServer):
    """ Fetch the status of the FTS server from RSS """

    res = self._resourceStatus.getElementStatus(ftsServer, 'FTS')
    if not res['OK']:
      return res

    result = res['Value']
    if ftsServer not in result:
      return S_ERROR("No FTS Server %s known to RSS" % ftsServer)

    if result[ftsServer]['all'] == 'Active':
      return S_OK(True)

    return S_OK(False)

  def chooseFTS3Server(self):
    """
      Choose the appropriate FTS3 server depending on the policy
    """

    fts3Server = None
    attempt = 0

    while not fts3Server and attempt < self._maxAttempts:

      fts3Server = self._policyMethod(attempt)
      res = self._getFTSServerStatus(fts3Server)

      if not res['OK']:
        self.log.warn("Error getting the RSS status for %s: %s" % (fts3Server, res))
        fts3Server = None
        attempt += 1
        continue

      ftsServerStatus = res['Value']

      if not ftsServerStatus:
        self.log.warn('FTS server %s is not in good shape. Choose another one' % fts3Server)
        fts3Server = None
        attempt += 1

    if fts3Server:
      return S_OK(self._serverDict[fts3Server])

    return S_ERROR("Could not find an FTS3 server (max attempt reached)")
