#!/usr/bin/env python
########################################################################
# :file:    dirac-dms-directory-sync
# :author:  Marko Petric
########################################################################
"""
Provides basic rsync functionality for DIRAC

Syncs the source destination folder recursivly into the target destination

If option --sync is used contend that is not in the source directory but is
only in the target directory will be deleted.

Usage:
  dirac-dms-directory-sync Source Destination

  e.g.: Download
    dirac-dms-directory-sync LFN Path
  or Upload
    dirac-dms-directory-sync Path LFN SE

Arguments:
  LFN:           Logical File Name (Path to directory)
  Path:          Local path to the file (Path to directory)
  SE:            DIRAC Storage Element
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import DIRAC
from DIRAC.Core.Base import Script
from DIRAC.Core.Utilities.DIRACScript import DIRACScript

__RCSID__ = "$Id$"


@DIRACScript()
def main():
  Script.registerSwitch("D", "sync", "Make target directory identical to source")
  Script.registerSwitch("j:", "parallel=", "Multithreaded download and upload")
  Script.parseCommandLine(ignoreErrors=False)

  args = Script.getPositionalArgs()
  if len(args) < 1 or len(args) > 3:
    Script.showHelp()

  sync = False
  parallel = 1
  for switch in Script.getUnprocessedSwitches():
    if switch[0].lower() == "s" or switch[0].lower() == "sync":
      sync = True
    if switch[0].lower() == "j" or switch[0].lower() == "parallel":
      parallel = int(switch[1])

  from DIRAC import S_OK, S_ERROR
  from DIRAC import gConfig, gLogger
  from DIRAC.Core.Utilities.List import breakListIntoChunks
  from DIRAC.Core.Utilities.ReturnValues import returnSingleResult
  from DIRAC.Core.Utilities.File import mkDir
  from DIRAC.Resources.Catalog.FileCatalog import FileCatalog
  from DIRAC.DataManagementSystem.Client.DataManager import DataManager
  from DIRAC.Resources.Storage.StorageElement import StorageElement

  def getSetOfLocalDirectoriesAndFiles(path):
    """Return a set of all directories and subdirectories and a set of
    files contained therein for a given local path
    """

    fullPath = os.path.abspath(path)

    if not os.path.isdir(fullPath):
      return S_ERROR("The path: " + fullPath + " does not exist!")

    directories = set()
    files = set()

    for dirname, dirnames, filenames in os.walk(path):
      # add path to all subdirectories first.
      for subdirname in dirnames:
        fullSubdirname = os.path.join(dirname, subdirname)
        fullSubdirname = os.path.abspath(fullSubdirname)
        fullSubdirname = fullSubdirname.replace(fullPath, '').lstrip('/')
        directories.add(fullSubdirname)
    # add path to all filenames.
      for filename in filenames:
        fullFilename = os.path.join(dirname, filename)
        fullFilename = os.path.abspath(fullFilename)
        fullFilename = fullFilename.replace(fullPath, '').lstrip('/')
        fileSize = os.path.getsize(fullPath + "/" + fullFilename)
        if fileSize > 0:
          files.add((fullFilename, int(fileSize)))

    tree = {}
    tree["Directories"] = directories
    tree["Files"] = files

    return S_OK(tree)

  def getSetOfRemoteSubDirectoriesAndFiles(path, fc, directories, files):
    """
    Recursively traverses all the subdirectories of a directory and returns a set of directories and files
    """
    result = fc.listDirectory(path)
    if result['OK']:
      if result['Value']['Successful']:
        for entry in result['Value']['Successful'][path]['Files']:
          size = result['Value']['Successful'][path]['Files'][entry]['MetaData']['Size']
          files.add((entry, size))
        for entry in result['Value']['Successful'][path]['SubDirs']:
          directories.add(entry)
          res = getSetOfRemoteSubDirectoriesAndFiles(entry, fc, directories, files)
          if not res['OK']:
            return S_ERROR('Error: ' + res['Message'])
        return S_OK()
      else:
        return S_ERROR("Error: %s" % result['Value'])
    else:
      return S_ERROR("Error:" + result['Message'])

  def getSetOfRemoteDirectoriesAndFiles(fc, path):
    """
    Return a set of all directories and subdirectories and the therein contained files for a given LFN
    """
    directories = set()
    files = set()

    res = getSetOfRemoteSubDirectoriesAndFiles(path, fc, directories, files)
    if not res['OK']:
      return S_ERROR('Could not list remote directory: ' + res['Message'])

    return_directories = set()
    return_files = set()

    for myfile in files:
      return_files.add((myfile[0].replace(path, '').lstrip('/'), myfile[1]))

    for mydirectory in directories:
      return_directories.add(mydirectory.replace(path, '').lstrip('/'))

    tree = {}
    tree["Directories"] = return_directories
    tree["Files"] = return_files

    return S_OK(tree)

  def isInFileCatalog(fc, path):
    """
    Check if the file is in the File Catalog
    """

    result = fc.listDirectory(path)
    if result['OK']:
      if result['Value']['Successful']:
        return S_OK()
      else:
        return S_ERROR()
    else:
      return S_ERROR()

  def getContentToSync(upload, fc, source_dir, dest_dir):
    """
    Return list of files and directories to be create and deleted
    """

    if upload:
      res = getSetOfRemoteDirectoriesAndFiles(fc, dest_dir)
      if not res['OK']:
        return S_ERROR(res['Message'])
      to_dirs = res['Value']['Directories']
      to_files = res['Value']['Files']

      res = getSetOfLocalDirectoriesAndFiles(source_dir)
      if not res['OK']:
        return S_ERROR(res['Message'])
      from_dirs = res['Value']['Directories']
      from_files = res['Value']['Files']

    else:
      res = getSetOfLocalDirectoriesAndFiles(dest_dir)
      if not res['OK']:
        return S_ERROR(res['Message'])
      to_dirs = res['Value']['Directories']
      to_files = res['Value']['Files']

      res = getSetOfRemoteDirectoriesAndFiles(fc, source_dir)
      if not res['OK']:
        return S_ERROR(res['Message'])
      from_dirs = res['Value']['Directories']
      from_files = res['Value']['Files']

    # Create list of directories to delete
    dirs_delete = list(to_dirs - from_dirs)
    # Sort the list by depth of directory tree
    dirs_delete.sort(key=lambda s: -s.count('/'))
    # Create list of directories to create
    dirs_create = list(from_dirs - to_dirs)
    # Sort the list by depth of directory tree
    dirs_create.sort(key=lambda s: s.count('/'))

    # Flatten the list of pairs (filename, size) to list of filename
    files_delete = [pair[0] for pair in list(to_files - from_files)]
    files_create = [pair[0] for pair in list(from_files - to_files)]

    create = {}
    create["Directories"] = dirs_create
    create["Files"] = files_create

    delete = {}
    delete["Directories"] = dirs_delete
    delete["Files"] = files_delete

    tree = {}
    tree["Create"] = create
    tree["Delete"] = delete

    return S_OK(tree)

  def removeRemoteFiles(dm, lfns):
    """
    Remove file from the catalog
    """
    for lfnList in breakListIntoChunks(lfns, 100):
      res = dm.removeFile(lfnList)
      if not res['OK']:
        return S_ERROR("Failed to remove files:" + lfnList + res['Message'])
      else:
        return S_OK()

  def uploadLocalFile(dm, lfn, localfile, storage):
    """
    Upload a local file to a storage element
    """
    res = dm.putAndRegister(lfn, localfile, storage, None)
    if not res['OK']:
      return S_ERROR('Error: failed to upload %s to %s' % (lfn, storage))
    else:
      return S_OK('Successfully uploaded file to %s' % storage)

  def downloadRemoteFile(dm, lfn, destination):
    """
    Download a file from the system
    """
    res = dm.getFile(lfn, destination)
    if not res['OK']:
      return S_ERROR('Error: failed to download %s ' % lfn)
    else:
      return S_OK('Successfully uploaded file %s' % lfn)

  def removeStorageDirectoryFromSE(directory, storageElement):
    """
    Delete directory on selected storage element
    """

    se = StorageElement(storageElement, False)
    res = returnSingleResult(se.exists(directory))

    if not res['OK']:
      return S_ERROR("Failed to obtain existence of directory" + res['Message'])

    exists = res['Value']
    if not exists:
      return S_OK("The directory %s does not exist at %s " % (directory, storageElement))

    res = returnSingleResult(se.removeDirectory(directory, recursive=True))
    if not res['OK']:
      return S_ERROR("Failed to remove storage directory" + res['Message'])

    return S_OK()

  def removeRemoteDirectory(fc, lfn):
    """
    Remove file from the catalog
    """
    storageElements = gConfig.getValue('Resources/StorageElementGroups/SE_Cleaning_List', [])

    for storageElement in sorted(storageElements):
      res = removeStorageDirectoryFromSE(lfn, storageElement)
      if not res['OK']:
        return S_ERROR("Failed to clean storage directory at all SE:" + res['Message'])
    res = returnSingleResult(fc.removeDirectory(lfn, recursive=True))
    if not res['OK']:
      return S_ERROR("Failed to clean storage directory at all SE:" + res['Message'])

    return S_OK("Successfully removed directory")

  def createRemoteDirectory(fc, newdir):
    """
    Create directory in file catalog
    """
    result = fc.createDirectory(newdir)
    if result['OK']:
      if result['Value']['Successful'] and newdir in result['Value']['Successful']:
        return S_OK("Successfully created directory:" + newdir)
      elif result['Value']['Failed'] and newdir in result['Value']['Failed']:
        return S_ERROR('Failed to create directory: ' + result['Value']['Failed'][newdir])
    else:
      return S_ERROR('Failed to create directory:' + result['Message'])

  def createLocalDirectory(directory):
    """
    Create local directory
    """
    mkDir(directory)
    if not os.path.exists(directory):
      return S_ERROR('Directory creation failed')
    return S_OK('Created directory successfully')

  def removeLocalFile(path):
    """
    Remove local file
    """
    try:
      os.remove(path)
    except OSError as e:
      return S_ERROR('Directory creation failed:' + e.strerror)

    if os.path.isfile(path):
      return S_ERROR('File deleting failed')
    return S_OK('Removed file successfully')

  def removeLocaDirectory(path):
    """
    Remove local directory
    """
    try:
      os.rmdir(path)
    except OSError as e:
      return S_ERROR('Deleting directory failed: ' + e.strerror)

    if os.path.isdir(path):
      return S_ERROR('Directory deleting failed')
    return S_OK('Removed directory successfully')

  def doUpload(fc, dm, result, source_dir, dest_dir, storage, delete, nthreads):
    """
    Wrapper for uploading files
    """
    if delete:
      lfns = [dest_dir + "/" + filename for filename in result['Value']['Delete']['Files']]
      if len(lfns) > 0:
        res = removeRemoteFiles(dm, lfns)
        if not res['OK']:
          gLogger.fatal('Deleting of files: ' + lfns + " -X- [FAILED]" + res['Message'])
          DIRAC.exit(1)
        else:
          gLogger.notice("Deleting " + ', '.join(lfns) + " -> [DONE]")

      for directoryname in result['Value']['Delete']['Directories']:
        res = removeRemoteDirectory(fc, dest_dir + "/" + directoryname)
        if not res['OK']:
          gLogger.fatal('Deleting of directory: ' + directoryname + " -X- [FAILED] " + res['Message'])
          DIRAC.exit(1)
        else:
          gLogger.notice("Deleting " + directoryname + " -> [DONE]")

    for directoryname in result['Value']['Create']['Directories']:
      res = createRemoteDirectory(fc, dest_dir + "/" + directoryname)
      if not res['OK']:
        gLogger.fatal('Creation of directory: ' + directoryname + " -X- [FAILED] " + res['Message'])
        DIRAC.exit(1)
      else:
        gLogger.notice("Creating " + directoryname + " -> [DONE]")

    listOfFiles = result['Value']['Create']['Files']
    # Chech that we do not have to many threads
    if nthreads > len(listOfFiles):
      nthreads = len(listOfFiles)

    if nthreads == 0:
      return S_OK('Upload finished successfully')

    listOfListOfFiles = chunkList(listOfFiles, nthreads)
    res = runInParallel(
        arguments=[
            dm,
            source_dir,
            dest_dir,
            storage],
        listOfLists=listOfListOfFiles,
        function=uploadListOfFiles)
    if not res['OK']:
      return S_ERROR("Upload of files failed")

    return S_OK('Upload finished successfully')

  def uploadListOfFiles(dm, source_dir, dest_dir, storage, listOfFiles, tID):
    """
    Wrapper for multithreaded uploading of a list of files
    """
    log = gLogger.getSubLogger("[Thread %s] " % tID)
    threadLine = "[Thread %s]" % tID
    for filename in listOfFiles:
      res = uploadLocalFile(dm, dest_dir + "/" + filename, source_dir + "/" + filename, storage)
      if not res['OK']:
        log.fatal(threadLine + ' Uploading ' + filename + ' -X- [FAILED] ' + res['Message'])
        DIRAC.exit(1)
      else:
        log.notice(threadLine + " Uploading " + filename + " -> [DONE]")

  def doDownload(dm, result, source_dir, dest_dir, delete, nthreads):
    """
    Wrapper for downloading files
    """
    if delete:
      for filename in result['Value']['Delete']['Files']:
        res = removeLocalFile(dest_dir + "/" + filename)
        if not res['OK']:
          gLogger.fatal('Deleting of file: ' + filename + ' -X- [FAILED] ' + res['Message'])
          DIRAC.exit(1)
        else:
          gLogger.notice("Deleting " + filename + " -> [DONE]")

      for directoryname in result['Value']['Delete']['Directories']:
        res = removeLocaDirectory(dest_dir + "/" + directoryname)
        if not res['OK']:
          gLogger.fatal('Deleting of directory: ' + directoryname + ' -X- [FAILED] ' + res['Message'])
          DIRAC.exit(1)
        else:
          gLogger.notice("Deleting " + directoryname + " -> [DONE]")

    for directoryname in result['Value']['Create']['Directories']:
      res = createLocalDirectory(dest_dir + "/" + directoryname)
      if not res['OK']:
        gLogger.fatal('Creation of directory: ' + directoryname + ' -X- [FAILED] ' + res['Message'])
        DIRAC.exit(1)
      else:
        gLogger.notice("Creating " + directoryname + " -> [DONE]")

    listOfFiles = result['Value']['Create']['Files']
    # Chech that we do not have to many threads
    if nthreads > len(listOfFiles):
      nthreads = len(listOfFiles)

    if nthreads == 0:
      return S_OK('Upload finished successfully')

    listOfListOfFiles = chunkList(listOfFiles, nthreads)
    res = runInParallel(
        arguments=[dm, source_dir, dest_dir],
        listOfLists=listOfListOfFiles,
        function=downloadListOfFiles,
    )

    if not res['OK']:
      return S_ERROR("Download of files failed")

    return S_OK('Upload finished successfully')

  def chunkList(alist, nchunks):
    """
    Split a list into a list of equaliy sized lists
    """
    avg = len(alist) / float(nchunks)
    out = []
    last = 0.0

    while last < len(alist):
      out.append(alist[int(last):int(last + avg)])
      last += avg

    return out

  def downloadListOfFiles(dm, source_dir, dest_dir, listOfFiles, tID):
    """
    Wrapper for multithreaded downloading of a list of files
    """
    log = gLogger.getSubLogger("[Thread %s] " % tID)
    threadLine = "[Thread %s]" % tID
    for filename in listOfFiles:
      res = downloadRemoteFile(dm, source_dir + "/" + filename, dest_dir + ("/" + filename).rsplit("/", 1)[0])
      if not res['OK']:
        log.fatal(threadLine + ' Downloading ' + filename + ' -X- [FAILED] ' + res['Message'])
        DIRAC.exit(1)
      else:
        log.notice(threadLine + " Downloading " + filename + " -> [DONE]")

  def runInParallel(arguments, listOfLists, function):
    """
    Helper for execution of uploads and downloads in parallel
    """
    from multiprocessing import Process
    processes = []
    for tID, alist in enumerate(listOfLists):
      argums = arguments + [alist] + [tID]
      pro = Process(target=function, args=argums)
      pro.start()
      processes.append(pro)
    for process in processes:
      process.join()

    for process in processes:
      if process.exitcode == 1:
        return S_ERROR()
    return S_OK()

  def syncDestinations(upload, source_dir, dest_dir, storage, delete, nthreads):
    """
    Top level wrapper to execute functions
    """

    fc = FileCatalog()
    dm = DataManager()

    result = getContentToSync(upload, fc, source_dir, dest_dir)
    if not result['OK']:
      return S_ERROR(result['Message'])

    if upload:
      res = doUpload(fc, dm, result, source_dir, dest_dir, storage, delete, nthreads)
      if not res['OK']:
        return S_ERROR('Upload failed: ' + res['Message'])
    else:
      res = doDownload(dm, result, source_dir, dest_dir, delete, nthreads)
      if not res['OK']:
        return S_ERROR('Download failed: ' + res['Message'])

    return S_OK('Mirroring successfully finished')

  def run(parameters, delete, nthreads):
    """
    The main user interface
    """

    source_dir = parameters[0]
    dest_dir = parameters[1]
    upload = False
    storage = None

    if len(parameters) == 3:
      storage = parameters[2]
      source_dir = os.path.abspath(source_dir)
      dest_dir = dest_dir.rstrip('/')
      upload = True
      if not os.path.isdir(source_dir):
        gLogger.fatal("Source directory does not exist")
        DIRAC.exit(1)

    if len(parameters) == 2:
      dest_dir = os.path.abspath(dest_dir)
      source_dir = source_dir.rstrip('/')
      if not os.path.isdir(dest_dir):
        gLogger.fatal("Destination directory does not exist")
        DIRAC.exit(1)

    res = syncDestinations(upload, source_dir, dest_dir, storage, delete, nthreads)
    if not res['OK']:
      return S_ERROR(res['Message'])

    return S_OK("Successfully mirrored " + source_dir + " into " + dest_dir)

  returnValue = run(args, sync, parallel)
  if not returnValue['OK']:
    gLogger.fatal(returnValue['Message'])
    DIRAC.exit(1)
  else:
    gLogger.notice(returnValue['Value'])
    DIRAC.exit(0)


if __name__ == "__main__":
  main()
