#!/usr/bin/env python
# -*- coding: utf-8 -*-

# -------------------------------------------------------------------
# Copyright (c) 2010-2019 Denis Machard
# This file is part of the extensive automation project
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
# MA 02110-1301 USA
# -------------------------------------------------------------------

"""
Module to construct messages
"""

try:
    import cPickle
except ImportError: # support python 3
    import pickle as cPickle
import zlib
import base64
import sys
import json
import time


# unicode = str with python3
if sys.version_info > (3,):
    unicode = str
    
HEAD_SEP            = b" "
BODY_SEP            = b"\n"

RSQ_CMD             = b"RSQ"
RSQ_NOTIFY          = b"NOTIFY"

RSP_CODE_OK         = ( b"200", b"OK" )
RSP_CODE_FORBIDDEN  = ( b"403", b"FORBIDDEN" )
RSP_CODE_FAILED     = ( b"400", b"FAILED" )
RSP_CODE_ERROR      = ( b"500", b"ERROR" )

CMD_ERROR           = -1
CMD_HELLO           = 0
CMD_GET_PROBE       = 1
CMD_START_PROBE     = 2
CMD_STOP_PROBE      = 3
CMD_NEW_FILE        = 4
CMD_INTERACT        = 5
CMD_START_AGENT     = 6
CMD_STOP_AGENT      = 7

def bytes2str(val):
    """
    bytes 2 str conversion, only for python3
    """
    if isinstance(val, bytes):
        return str(val, "utf8")
    else:
        return val
        
        
import pickle


def bytes_to_unicode(ob):
    """
    Byte to unicode with exception...
    Only for py3, will be removed on future version...
    """
    t = type(ob)
    if t in (list, tuple):
        try:
            l = [str(i, 'utf-8') if type(i) is bytes else i for i in ob]
        except UnicodeDecodeError as e:
            l = [ i for i in ob] # keep as bytes
        l = [bytes_to_unicode(i) if type(i) in (list, tuple, dict) else i for i in l]
        ro = tuple(l) if t is tuple else l
    elif t is dict:
        byte_keys = [i for i in ob if type(i) is bytes]
        for bk in byte_keys:
            v = ob[bk]
            del(ob[bk])
            try:
                ob[str(bk,'utf-8')] = v
            except UnicodeDecodeError as e:
                ob[bk] = v # keep as bytes
        for k in ob:
            if type(ob[k]) is bytes:
                try:
                    ob[k] = str(ob[k], 'utf-8')
                except UnicodeDecodeError as e:
                    ob[k] = ob[k] # keep as bytes
            elif type(ob[k]) in (list, tuple, dict):
                ob[k] = bytes_to_unicode(ob[k])
        ro = ob
    else:
        ro = ob
        #print("unprocessed object: {0} {1}".format(t, ob))
    return ro
    
class Messages(object):
    """
    Codec for messages
    """
    def __init__ (self, userId = None, useJson=False, pickleVer=2):
        """
        Constructor

        @param userId: user identifier (default=None)
        @type userId: string
        
        @param useJson: use json instead of pickle (default=False)
        @type useJson: boolean
        
        @param pickleVer: pickle version used (default=2) backward compatibility with python2
        @type pickleVer: integer
        """
        self.__userId = userId
        self.__useJson = useJson
        self.__pickleProtocol = pickleVer

    def setUserId (self, userId):
        """
        Set the userid

        @param userId: user identifier
        @type userId: string
        """
        self.__userId = userId

    def encode (self, op, tid, desc, body = None):
        """
        Encode a message (response or request

        @param op: operator
        @type op:

        @param tid: transaction id
        @type tid:

        @param desc: description
        @type desc:

        @param body: content of the message
        @type body:

        @return: the message encoded
        @rtype: string
        """
        # encode a response ?
        if op.isdigit():
            # example: 200 OK, 403 FORBIDDEN, etc...
            if op not in [ RSP_CODE_OK[0], RSP_CODE_FORBIDDEN[0], RSP_CODE_FAILED[0], RSP_CODE_ERROR[0] ]:
                raise Exception('unknown response code: %s' % op)
            if desc not in [ RSP_CODE_OK[1], RSP_CODE_FORBIDDEN[1], RSP_CODE_FAILED[1], RSP_CODE_ERROR[1] ]:
                raise Exception('unknown response phrase: %s' % desc)
                
        # encode a request: cmd or notify
        else:
            if op not in [ RSQ_CMD, RSQ_NOTIFY ]:
                raise Exception('unknown request: %s' % op)

        # prepare the head of the message
        if sys.version_info > (3,): # python 3 support
            if not isinstance(desc, bytes):
                desc = bytes(desc, 'utf8')
            ret = [ HEAD_SEP.join( [ op, bytes(str(tid), 'utf8'), desc ] ) ]
        else:
            ret = [ HEAD_SEP.join( [ op, str(tid), desc ] ) ]

        # adding the body as json or pickled(default)
        # also zipped and encoded in base64
        if body: 
            if not self.__useJson:
                pickled = cPickle.dumps(body, protocol=self.__pickleProtocol)
                bod = zlib.compress(pickled)
                ret.append( base64.encodestring(bod) )
            else:
                json_data = json.dumps(body, ensure_ascii=False)
                compressed = zlib.compress( json_data )
                ret.append( base64.encodestring(compressed) )
        else: 
            rslt = ret

        # join all and return the message
        rslt = BODY_SEP.join(ret)

        return rslt

    def decode (self, msgraw):
        """
        Decode a message: request or response

        @param msgraw:
        @type msgraw:

        @return: request or response
        @rtype: typle
        """
        # split the body
        msg = msgraw.split(BODY_SEP, 1)
        # no body in the message
        if len(msg) == 1:
            head = msg[0].split(HEAD_SEP)
            body = b''
            
        # header with body
        elif len(msg) == 2:
            # extract the header of the message
            head = msg[0].split(HEAD_SEP)
            
            # extract the body 
            # the body is encoded in the base 64 and zipped
            if sys.version_info > (3,): # support python3
                decoded = base64.b64decode(msg[1])
            else:
                decoded = base64.decodestring(msg[1]) # deprecated function
            decompressed_data = zlib.decompress(decoded)
            if not self.__useJson:
                if sys.version_info > (3,): # support python3
                    body =  cPickle.loads( decompressed_data, encoding="bytes")
                    body = bytes_to_unicode(body) # convert bytes to unicode with exceptions, workaround
                else:
                    body =  cPickle.loads( decompressed_data )
            else:
                body = json.loads( decompressed_data, encoding="ISO-8859-1" )
        else:
            raise Exception('invalid message')

        # the final message is a dictionary with the following keys
        #  - tid key = transaction id
        #  - body key = content of the message
        # if response
        #  - code key
        #  - phrase key
        # if request
        #  - cmd key
        #  - userid key
        # finally return a tuple (response,...) or (request, ...)
        ret = {}
        ret['tid'] = int(head[1])
        ret['body'] = body
        if head[0].isdigit():
            ret['code'] = head[0]
            if ret['code'] not in [ RSP_CODE_OK[0], RSP_CODE_FORBIDDEN[0], RSP_CODE_FAILED[0], RSP_CODE_ERROR[0] ]:
                raise Exception('unknown response code: %s' % ret['code'])

            ret['phrase'] = head[2]
            if ret['phrase'] not in [ RSP_CODE_OK[1], RSP_CODE_FORBIDDEN[1], RSP_CODE_FAILED[1], RSP_CODE_ERROR[1] ]:
                raise Exception('unknown response phrase: %s' % ret['phrase'] )
            return ('response', ret)
        else:

            ret['cmd'] = head[0]
            if ret['cmd'] not in [ RSQ_CMD, RSQ_NOTIFY ]:
                raise Exception('unknown request: %s' % ret['cmd'])
            ret['userid'] = head[2]
            return ('request', ret)
    
    def ok ( self, tid, body = None):   
        """
        Encode a ok message

        @param tid: transaction id
        @type tid:

        @param body:
        @type body:

        @return: one ok message encoded
        @rtype: string
        """
        ret = self.encode(op = RSP_CODE_OK[0], tid = tid, desc = RSP_CODE_OK[1], body = body)
        return ret

    def forbidden ( self, tid, body = None):
        """
        Encode a forbidden message

        @param tid: transaction id
        @type tid:

        @param body:
        @type body:

        @return: one forbidden message encoded
        @rtype: string
        """
        ret = self.encode(op = RSP_CODE_FORBIDDEN[0], tid = tid, desc = RSP_CODE_FORBIDDEN[1], body = body)
        return ret

    def failed ( self, tid, body = None):
        """
        Encode a failed message

        @param tid: transaction id
        @type tid:

        @param body:
        @type body:

        @return: one failed message encoded
        @rtype: string
        """
        ret = self.encode(op = RSP_CODE_FAILED[0], tid = tid, desc = RSP_CODE_FAILED[1], body = body)
        return ret

    def error ( self, tid, body = None):
        """
        Encode an error message

        @param tid: transaction id
        @type tid:

        @param body:
        @type body:

        @return: one error message encoded
        @rtype: string
        """
        ret = self.encode(op = RSP_CODE_ERROR[0], tid = tid, desc = RSP_CODE_ERROR[1], body = body)
        return ret

    def cmd (self, tid, body):
        """
        Encode a command

        @param tid: transaction id
        @type tid:

        @param body:
        @type body:

        @return: one cmd message encoded
        @rtype: string
        """
        ret = self.encode(op = RSQ_CMD, tid = tid, desc = self.__userId, body = body)
        return ret

    def notify (self, tid, body):
        """
        Encode a notify

        @param tid: transaction id
        @type tid:

        @param body:
        @type body:

        @return: one notify message encoded
        @rtype: string
        """
        ret = self.encode(op = RSQ_NOTIFY, tid = tid, desc = self.__userId, body = body)
        return ret
