from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
__RCSID__ = "$Id$"

import hashlib
import io
import os
import tarfile
import tempfile
import threading

import six
from six import StringIO

from DIRAC.Core.Utilities.ReturnValues import S_OK, S_ERROR
from DIRAC.FrameworkSystem.Client.Logger import gLogger

try:
    # Python 2: "file" is built-in
    file_types = file, io.IOBase
except NameError:
    # Python 3: "file" fully replaced with IOBase
    file_types = (io.IOBase,)

gLogger = gLogger.getSubLogger("FileTransmissionHelper")


class FileHelper(object):

  __validDirections = ("toClient", "fromClient", 'receive', 'send')
  __directionsMapping = {'toClient': 'send', 'fromClient': 'receive'}

  def __init__(self, oTransport=None, checkSum=True):
    self.oTransport = oTransport
    self.__checkMD5 = checkSum
    self.__oMD5 = hashlib.md5()
    self.bFinishedTransmission = False
    self.bReceivedEOF = False
    self.direction = False
    self.packetSize = 1048576
    self.__fileBytes = 0
    self.__log = gLogger.getSubLogger("FileHelper")

  def disableCheckSum(self):
    self.__checkMD5 = False

  def enableCheckSum(self):
    self.__checkMD5 = True

  def setTransport(self, oTransport):
    self.oTransport = oTransport

  def setDirection(self, direction):
    if direction in FileHelper.__validDirections:
      if direction in FileHelper.__directionsMapping:
        self.direction = FileHelper.__directionsMapping[direction]
      else:
        self.direction = direction

  def getHash(self):
    return self.__oMD5.hexdigest()

  def getTransferedBytes(self):
    return self.__fileBytes

  def sendData(self, sBuffer):
    if six.PY3 and isinstance(sBuffer, str):
      sBuffer = sBuffer.encode(errors="surrogateescape")
    if self.__checkMD5:
      self.__oMD5.update(sBuffer)
    retVal = self.oTransport.sendData(S_OK([True, sBuffer]))
    if not retVal['OK']:
      return retVal
    retVal = self.oTransport.receiveData()
    return retVal

  def sendEOF(self):
    retVal = self.oTransport.sendData(S_OK([False, self.__oMD5.hexdigest()]))
    if not retVal['OK']:
      return retVal
    self.__finishedTransmission()
    return S_OK()

  def sendError(self, errorMsg):
    retVal = self.oTransport.sendData(S_ERROR(errorMsg))
    if not retVal['OK']:
      return retVal
    self.__finishedTransmission()
    return S_OK()

  def receiveData(self, maxBufferSize=0):
    retVal = self.oTransport.receiveData(maxBufferSize=maxBufferSize)
    if 'AbortTransfer' in retVal and retVal['AbortTransfer']:
      self.oTransport.sendData(S_OK())
      self.__finishedTransmission()
      self.bReceivedEOF = True
      return S_OK('')
    if not retVal['OK']:
      return retVal
    stBuffer = retVal['Value']
    if stBuffer[0]:
      if six.PY3 and isinstance(stBuffer[1], str):
        stBuffer[1] = stBuffer[1].encode(errors="surrogateescape")
      if self.__checkMD5:
        self.__oMD5.update(stBuffer[1])
      self.oTransport.sendData(S_OK())
    else:
      self.bReceivedEOF = True
      if self.__checkMD5 and not self.__oMD5.hexdigest() == stBuffer[1]:
        self.bErrorInMD5 = True
      self.__finishedTransmission()
      return S_OK("")
    return S_OK(stBuffer[1])

  def receivedEOF(self):
    return self.bReceivedEOF

  def markAsTransferred(self):
    if not self.bFinishedTransmission:
      if self.direction == "receive":
        self.oTransport.receiveData()
        abortTrans = S_OK()
        abortTrans['AbortTransfer'] = True
        self.oTransport.sendData(abortTrans)
      else:
        abortTrans = S_OK([False, ""])
        abortTrans['AbortTransfer'] = True
        retVal = self.oTransport.sendData(abortTrans)
        if not retVal['OK']:
          return retVal
        self.oTransport.receiveData()
    self.__finishedTransmission()

  def __finishedTransmission(self):
    self.bFinishedTransmission = True

  def finishedTransmission(self):
    return self.bFinishedTransmission

  def errorInTransmission(self):
    return self.bErrorInMD5

  def networkToString(self, maxFileSize=0):
    """ Receive the input from a DISET client and return it as a string
    """

    stringIO = StringIO()
    result = self.networkToDataSink(stringIO, maxFileSize=maxFileSize)
    if not result['OK']:
      return result
    return S_OK(stringIO.getvalue())

  def networkToFD(self, iFD, maxFileSize=0):
    dataSink = os.fdopen(iFD, "w")
    try:
      return self.networkToDataSink(dataSink, maxFileSize=maxFileSize)
    finally:
      try:
        dataSink.close()
      except Exception as e:
        pass

  def networkToDataSink(self, dataSink, maxFileSize=0):
    if "write" not in dir(dataSink):
      return S_ERROR("%s data sink object does not have a write method" % str(dataSink))
    self.__oMD5 = hashlib.md5()
    self.bReceivedEOF = False
    self.bErrorInMD5 = False
    receivedBytes = 0
    # try:
    result = self.receiveData(maxBufferSize=maxFileSize)
    if not result['OK']:
      return result
    strBuffer = result['Value']
    if six.PY3 and isinstance(strBuffer, str):
      strBuffer = strBuffer.encode(errors="surrogateescape")
    receivedBytes += len(strBuffer)
    while not self.receivedEOF():
      if maxFileSize > 0 and receivedBytes > maxFileSize:
        self.sendError("Exceeded maximum file size")
        return S_ERROR("Received file exceeded maximum size of %s bytes" % (maxFileSize))
      dataSink.write(strBuffer)
      result = self.receiveData(maxBufferSize=(maxFileSize - len(strBuffer)))
      if not result['OK']:
        return result
      strBuffer = result['Value']
      if six.PY3 and isinstance(strBuffer, str):
        strBuffer = strBuffer.encode(errors="surrogateescape")
      receivedBytes += len(strBuffer)
    if strBuffer:
      dataSink.write(strBuffer)
    # except Exception as e:
    #   return S_ERROR("Error while receiving file, %s" % str(e))
    if self.errorInTransmission():
      return S_ERROR("Error in the file CRC")
    self.__fileBytes = receivedBytes
    return S_OK()

  def stringToNetwork(self, stringVal):
    """ Send a given string to the DISET client over the network
    """

    stringIO = StringIO(stringVal)

    iPacketSize = self.packetSize
    ioffset = 0
    strlen = len(stringVal)
    try:
      while (ioffset) < strlen:
        if (ioffset + iPacketSize) < strlen:
          result = self.sendData(stringVal[ioffset:ioffset + iPacketSize])
        else:
          result = self.sendData(stringVal[ioffset:strlen])
        if not result['OK']:
          return result
        if 'AbortTransfer' in result and result['AbortTransfer']:
          self.__log.verbose("Transfer aborted")
          return S_OK()
        ioffset += iPacketSize
      self.sendEOF()
    except Exception as e:
      return S_ERROR("Error while sending string: %s" % str(e))
    try:
      stringIO.close()
    except Exception:
      pass
    return S_OK()

  def FDToNetwork(self, iFD):
    self.__oMD5 = hashlib.md5()
    iPacketSize = self.packetSize
    self.__fileBytes = 0
    sentBytes = 0
    try:
      sBuffer = os.read(iFD, iPacketSize)
      while len(sBuffer) > 0:
        dRetVal = self.sendData(sBuffer)
        if not dRetVal['OK']:
          return dRetVal
        if 'AbortTransfer' in dRetVal and dRetVal['AbortTransfer']:
          self.__log.verbose("Transfer aborted")
          return S_OK()
        sentBytes += len(sBuffer)
        sBuffer = os.read(iFD, iPacketSize)
      self.sendEOF()
    except Exception as e:
      gLogger.exception("Error while sending file")
      return S_ERROR("Error while sending file: %s" % str(e))
    self.__fileBytes = sentBytes
    return S_OK()

  def BufferToNetwork(self, stringToSend):
    sIO = StringIO(stringToSend)
    try:
      return self.DataSourceToNetwork(sIO)
    finally:
      sIO.close()

  def DataSourceToNetwork(self, dataSource):
    if "read" not in dir(dataSource):
      return S_ERROR("%s data source object does not have a read method" % str(dataSource))
    self.__oMD5 = hashlib.md5()
    iPacketSize = self.packetSize
    self.__fileBytes = 0
    sentBytes = 0
    try:
      sBuffer = dataSource.read(iPacketSize)
      while len(sBuffer) > 0:
        dRetVal = self.sendData(sBuffer)
        if not dRetVal['OK']:
          return dRetVal
        if 'AbortTransfer' in dRetVal and dRetVal['AbortTransfer']:
          self.__log.verbose("Transfer aborted")
          return S_OK()
        sentBytes += len(sBuffer)
        sBuffer = dataSource.read(iPacketSize)
      self.sendEOF()
    except Exception as e:
      gLogger.exception("Error while sending file")
      return S_ERROR("Error while sending file: %s" % str(e))
    self.__fileBytes = sentBytes
    return S_OK()

  def getFileDescriptor(self, uFile, sFileMode):
    closeAfter = True
    if isinstance(uFile, six.string_types):
      try:
        self.oFile = open(uFile, sFileMode)
      except IOError:
        return S_ERROR("%s can't be opened" % uFile)
      iFD = self.oFile.fileno()
    elif isinstance(uFile, file_types):
      iFD = uFile.fileno()
    elif isinstance(uFile, int):
      iFD = uFile
      closeAfter = False
    else:
      return S_ERROR("%s is not a valid file." % uFile)
    result = S_OK(iFD)
    result['closeAfterUse'] = closeAfter
    return result

  def getDataSink(self, uFile):
    closeAfter = True
    if isinstance(uFile, six.string_types):
      try:
        oFile = open(uFile, "wb")
      except IOError:
        return S_ERROR("%s can't be opened" % uFile)
    elif isinstance(uFile, file_types):
      oFile = uFile
      closeAfter = False
    elif isinstance(uFile, int):
      oFile = os.fdopen(uFile, "wb")
      closeAfter = True
    elif "write" in dir(uFile):
      oFile = uFile
      closeAfter = False
    else:
      return S_ERROR("%s is not a valid file." % uFile)
    result = S_OK(oFile)
    result['closeAfterUse'] = closeAfter
    return result

  def __createTar(self, fileList, wPipe, compress, autoClose=True):
    if 'write' in dir(wPipe):
      filePipe = wPipe
    else:
      filePipe = os.fdopen(wPipe, "w")
    tarMode = "w|"
    if compress:
      tarMode = "w|bz2"

    with tarfile.open(name="Pipe", mode=tarMode, fileobj=filePipe) as tar:
      for entry in fileList:
        tar.add(os.path.realpath(entry), os.path.basename(entry), recursive=True)
    if autoClose:
      try:
        filePipe.close()
      except Exception:
        pass

  def bulkToNetwork(self, fileList, compress=True, onthefly=True):
    if not onthefly:
      try:
        filePipe, filePath = tempfile.mkstemp()
      except Exception as e:
        return S_ERROR("Can't create temporary file to pregenerate the bulk: %s" % str(e))
      self.__createTar(fileList, filePipe, compress)
      try:
        fo = open(filePath, 'rb')
      except Exception as e:
        return S_ERROR("Can't read pregenerated bulk: %s" % str(e))
      result = self.DataSourceToNetwork(fo)
      try:
        fo.close()
        os.unlink(filePath)
      except Exception:
        pass
      return result
    else:
      rPipe, wPipe = os.pipe()
      thrd = threading.Thread(target=self.__createTar, args=(fileList, wPipe, compress))
      thrd.start()
      response = self.FDToNetwork(rPipe)
      try:
        os.close(rPipe)
      except Exception:
        pass
      return response

  def __extractTar(self, destDir, rPipe, compress):
    filePipe = os.fdopen(rPipe, "r")
    tarMode = "r|*"
    if compress:
      tarMode = "r|bz2"
    with tarfile.open(mode=tarMode, fileobj=filePipe) as tar:
      for tarInfo in tar:
        tar.extract(tarInfo, destDir)
    try:
      filePipe.close()
    except Exception:
      pass

  def __receiveToPipe(self, wPipe, retList, maxFileSize):
    retList.append(self.networkToFD(wPipe, maxFileSize=maxFileSize))
    try:
      os.close(wPipe)
    except Exception:
      pass

  def networkToBulk(self, destDir, compress=True, maxFileSize=0):
    retList = []
    rPipe, wPipe = os.pipe()
    thrd = threading.Thread(target=self.__receiveToPipe, args=(wPipe, retList, maxFileSize))
    thrd.start()
    try:
      self.__extractTar(destDir, rPipe, compress)
    except Exception as e:
      return S_ERROR("Error while extracting bulk: %s" % e)
    thrd.join()
    return retList[0]

  def bulkListToNetwork(self, iFD, compress=True):
    filePipe = os.fdopen(iFD, "r")
    try:
      tarMode = "r|"
      if compress:
        tarMode = "r|bz2"
      entries = []
      with tarfile.open(mode=tarMode, fileobj=filePipe) as tar:
        for tarInfo in tar:
          entries.append(tarInfo.name)
      filePipe.close()
      return S_OK(entries)
    except tarfile.ReadError as v:
      return S_ERROR("Error reading bulk: %s" % str(v))
    except tarfile.CompressionError as v:
      return S_ERROR("Error in bulk compression setting: %s" % str(v))
    except Exception as v:
      return S_ERROR("Error in listing bulk: %s" % str(v))
