""" Hosts BaseTransport class, which is a base for PlainTransport and SSLTransport

BaseTransport is used by the client and the service, it describes an opened connection.
Here a diagram of basic Client/Service exchange

Client -> ServiceReactor : Connect

Client<-->Service        : Handshake (only in SSLTransport)

Client -> Service        : Propose action

Client <- Service        : S_OK

Client -> RequestHandler : Arguments

Client <- RequestHandler : Response

Client <- Service        : Close
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

__RCSID__ = "$Id$"

import time
import six
from six import BytesIO
from hashlib import md5

try:
    import selectors
except ImportError:
    import selectors2 as selectors

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


class BaseTransport(object):
    """Invokes MixedEncode for marshaling/unmarshaling of data calls in transit"""

    bAllowReuseAddress = True

    # This option corresponds to the `backlog` param of the `listen` syscall.
    # you may want to read the man page before tuning it...
    iListenQueueSize = 128
    iReadTimeout = 600
    keepAliveMagic = b"dka"

    def __init__(self, stServerAddress, bServerMode=False, **kwargs):
        self.bServerMode = bServerMode
        self.extraArgsDict = kwargs
        self.byteStream = b""
        self.packetSize = 1048576  # 1MiB
        self.stServerAddress = stServerAddress
        self.peerCredentials = {}
        self.remoteAddress = False
        self.appData = ""
        self.startedKeepAlives = set()
        self.keepAliveId = md5((str(stServerAddress) + str(bServerMode)).encode()).hexdigest()
        self.receivedMessages = []
        self.sentKeepAlives = 0
        self.waitingForKeepAlivePong = False
        self.__keepAliveLapse = 0
        self.oSocket = None
        if "keepAliveLapse" in kwargs:
            try:
                self.__keepAliveLapse = max(150, int(kwargs["keepAliveLapse"]))
            except Exception:
                pass
        self.iListenQueueSize = max(self.iListenQueueSize, int(kwargs.get("SocketBacklog", 0)))
        self.__lastActionTimestamp = time.time()
        self.__lastServerRenewTimestamp = self.__lastActionTimestamp

    def __updateLastActionTimestamp(self):
        self.__lastActionTimestamp = time.time()

    def getLastActionTimestamp(self):
        return self.__lastActionTimestamp

    def getKeepAliveLapse(self):
        return self.__keepAliveLapse

    def handshake(self):
        """This method is overwritten by SSLTransport if we use a secured transport."""
        return S_OK()

    def close(self):
        self.oSocket.close()

    def setAppData(self, appData):
        self.appData = appData

    def getAppData(self):
        return self.appData

    def renewServerContext(self):
        self.__lastServerRenewTimestamp = time.time()
        return S_OK()

    def latestServerRenewTime(self):
        return self.__lastServerRenewTimestamp

    def getConnectingCredentials(self):
        """

        :return: dictionnary with credentials

          Return empty dictionnary for plainTransport.

          In SSLTransport it contains (after the handshake):

           - 'DN' : All identity name, e.g. ```/C=ch/O=DIRAC/OU=DIRAC CI/CN=ciuser```
           - 'CN' : Only the user name e.g. ciuser
           - 'x509Chain' : List of all certificates in the chain
           - 'isProxy' : True if the client use proxy certificate
           - 'isLimitedProxy' : True if the client use limited proxy certificate
           - 'group' (optional): Dirac group attached to the client
           - 'extraCredentials' (optional): Extra credentials if exists

          Before the handshake, dictionary is empty
        """
        return self.peerCredentials

    def setExtraCredentials(self, extraCredentials):
        """Add extra credentials to peerCredentials

        :param extraCredentials: group or tuple with DN and group
        :type extraCredentials: str or tuple
        """
        self.peerCredentials["extraCredentials"] = extraCredentials

    def serverMode(self):
        return self.bServerMode

    def getRemoteAddress(self):
        return self.remoteAddress

    def getLocalAddress(self):
        return self.oSocket.getsockname()

    def getSocket(self):
        return self.oSocket

    def _readReady(self):
        if not self.iReadTimeout:
            return True
        sel = selectors.DefaultSelector()
        sel.register(self.oSocket, selectors.EVENT_READ)
        if sel.select(timeout=self.iReadTimeout):
            return True
        return False

    def _read(self, bufSize=4096, skipReadyCheck=False):
        try:
            if skipReadyCheck or self._readReady():
                data = self.oSocket.recv(bufSize)
                if not data:
                    return S_ERROR("Connection closed by peer")
                else:
                    return S_OK(data)
            else:
                return S_ERROR("Connection seems stalled. Closing...")
        except Exception as e:
            return S_ERROR("Exception while reading from peer: %s" % str(e))

    def _write(self, buf):
        return S_OK(self.oSocket.send(buf))

    def sendData(self, uData, prefix=b""):
        self.__updateLastActionTimestamp()
        sCodedData = MixedEncode.encode(uData)
        if isinstance(sCodedData, six.text_type):
            sCodedData = sCodedData.encode()
        dataToSend = b"".join([prefix, str(len(sCodedData)).encode(), b":", sCodedData])
        for index in range(0, len(dataToSend), self.packetSize):
            bytesToSend = min(self.packetSize, len(dataToSend) - index)
            packSentBytes = 0
            while packSentBytes < bytesToSend:
                try:
                    result = self._write(dataToSend[index + packSentBytes : index + bytesToSend])
                    if not result["OK"]:
                        return result
                    sentBytes = result["Value"]
                except Exception as e:
                    return S_ERROR("Exception while sending data: %s" % e)
                if sentBytes == 0:
                    return S_ERROR("Connection closed by peer")
                packSentBytes += sentBytes
        del sCodedData
        sCodedData = None
        return S_OK()

    def receiveData(self, maxBufferSize=0, blockAfterKeepAlive=True, idleReceive=False):
        self.__updateLastActionTimestamp()
        if self.receivedMessages:
            return self.receivedMessages.pop(0)
        # Buffer size can't be less than 0
        maxBufferSize = max(maxBufferSize, 0)
        try:
            # Look either for message length of keep alive magic string
            iSeparatorPosition = self.byteStream.find(b":", 0, 10)
            keepAliveMagicLen = len(BaseTransport.keepAliveMagic)
            isKeepAlive = self.byteStream.find(BaseTransport.keepAliveMagic, 0, keepAliveMagicLen) == 0
            # While not found the message length or the ka, keep receiving
            while iSeparatorPosition == -1 and not isKeepAlive:
                retVal = self._read(16384)
                # If error return
                if not retVal["OK"]:
                    return retVal
                # If closed return error
                if not retVal["Value"]:
                    return S_ERROR("Peer closed connection")
                # New data!
                self.byteStream += retVal["Value"]
                # Look again for either message length of ka magic string
                iSeparatorPosition = self.byteStream.find(b":", 0, 10)
                isKeepAlive = self.byteStream.find(BaseTransport.keepAliveMagic, 0, keepAliveMagicLen) == 0
                # Over the limit?
                if maxBufferSize and len(self.byteStream) > maxBufferSize and iSeparatorPosition == -1:
                    return S_ERROR("Read limit exceeded (%s chars)" % maxBufferSize)
            # Keep alive magic!
            if isKeepAlive:
                gLogger.debug("Received keep alive header")
                # Remove the ka magic from the buffer and process the keep alive
                self.byteStream = self.byteStream[keepAliveMagicLen:]
                return self.__processKeepAlive(maxBufferSize, blockAfterKeepAlive)
            # From here it must be a real message!
            # Process the size and remove the msg length from the bytestream
            pkgSize = int(self.byteStream[:iSeparatorPosition])
            pkgData = self.byteStream[iSeparatorPosition + 1 :]
            readSize = len(pkgData)
            if readSize >= pkgSize:
                # If we already have all the data we need
                data = pkgData[:pkgSize]
                self.byteStream = pkgData[pkgSize:]
            else:
                # If we still need to read stuff
                pkgMem = BytesIO()
                pkgMem.write(pkgData)
                # Receive while there's still data to be received
                while readSize < pkgSize:
                    retVal = self._read(pkgSize - readSize, skipReadyCheck=True)
                    if not retVal["OK"]:
                        return retVal
                    if not retVal["Value"]:
                        return S_ERROR("Peer closed connection")
                    rcvData = retVal["Value"]
                    readSize += len(rcvData)
                    pkgMem.write(rcvData)
                    if maxBufferSize and readSize > maxBufferSize:
                        return S_ERROR("Read limit exceeded (%s chars)" % maxBufferSize)
                # Data is here! take it out from the bytestream, dencode and return
                if readSize == pkgSize:
                    data = pkgMem.getvalue()
                    self.byteStream = b""
                else:  # readSize > pkgSize:
                    pkgMem.seek(0, 0)
                    data = pkgMem.read(pkgSize)
                    self.byteStream = pkgMem.read()
            try:
                data = MixedEncode.decode(data)[0]
            except Exception as e:
                return S_ERROR("Could not decode received data: %s" % str(e))
            if idleReceive:
                self.receivedMessages.append(data)
                return S_OK()
            return data
        except Exception as e:
            gLogger.exception("Network error while receiving data")
            return S_ERROR("Network error while receiving data: %s" % str(e))

    def __processKeepAlive(self, maxBufferSize, blockAfterKeepAlive=True):
        gLogger.debug("Received Keep Alive")
        # Next message down the stream will be the ka data
        result = self.receiveData(maxBufferSize, blockAfterKeepAlive=False)
        if not result["OK"]:
            gLogger.debug("Error while receiving keep alive: %s" % result["Message"])
            return result
        # Is it a valid ka?
        kaData = result["Value"]
        for reqField in ("id", "kaping"):
            if reqField not in kaData:
                errMsg = "Invalid keep alive, missing %s" % reqField
                gLogger.debug(errMsg)
                return S_ERROR(errMsg)
        gLogger.debug("Received keep alive id %s" % kaData)
        # Need to check if it's one of the keep alives we sent or one started from the other side
        if kaData["kaping"]:
            # This is a keep alive PING. Let's send the PONG
            self.sendKeepAlive(responseId=kaData["id"])
        else:
            # If it's a pong then we flag that we don't need to wait for a pong
            self.waitingForKeepAlivePong = False
        # No blockAfterKeepAlive means return without further read
        if not blockAfterKeepAlive:
            result = S_OK()
            result["keepAlive"] = True
            return result
        # Let's listen for the next message downstream
        return self.receiveData(maxBufferSize, blockAfterKeepAlive)

    def sendKeepAlive(self, responseId=None, now=False):
        # If not responseId or not keepAliveLapse or not enough time has passed don't send keep alive
        if not responseId:
            if not self.__keepAliveLapse:
                return S_OK()
            if not now:
                now = time.time()
            if now - self.__lastActionTimestamp < self.__keepAliveLapse:
                return S_OK()

        self.__updateLastActionTimestamp()
        if responseId:
            self.waitingForKeepAlivePong = False
            kaData = S_OK({"id": responseId, "kaping": False})
        else:
            if self.waitingForKeepAlivePong:
                return S_OK()
            idK = self.keepAliveId + str(self.sentKeepAlives)
            self.sentKeepAlives += 1
            kaData = S_OK({"id": idK, "kaping": True})
            self.waitingForKeepAlivePong = True
        return self.sendData(kaData, prefix=BaseTransport.keepAliveMagic)

    def getFormattedCredentials(self):
        peerCreds = self.getConnectingCredentials()
        address = self.getRemoteAddress()
        if "username" in peerCreds:
            peerId = "[%s:%s]" % (peerCreds["group"], peerCreds["username"])
        else:
            peerId = ""
        if address[0].find(":") > -1:
            return "([%s]:%s)%s" % (address[0], address[1], peerId)
        return "(%s:%s)%s" % (address[0], address[1], peerId)

    def setSocketTimeout(self, timeout):
        """
        This method has to be overwritten, if we want to increase the socket timeout.
        """
        pass
