import struct

class NMBError(Exception): pass

class NotConnectedError(NMBError):
    """
    Raisd when the underlying NMB connection has been disconnected or not connected yet
    """
    pass

class NMBSessionMessage:

    HEADER_STRUCT_FORMAT = '>BBH'
    HEADER_STRUCT_SIZE = struct.calcsize(HEADER_STRUCT_FORMAT)

    def __init__(self):
        self.reset()

    def reset(self):
        self.type = 0
        self.flags = 0
        self.data = ''

    def decode(self, data, offset):
        data_len = len(data)

        if data_len < offset + self.HEADER_STRUCT_SIZE:
            # Not enough data for decoding
            return 0

        self.reset()
        self.type, self.flags, length = struct.unpack(self.HEADER_STRUCT_FORMAT, data[offset:offset+self.HEADER_STRUCT_SIZE])

        if self.flags & 0x01:
            length |= 0x010000

        if data_len < offset + self.HEADER_STRUCT_SIZE + length:
            return 0

        self.data = data[offset+self.HEADER_STRUCT_SIZE:offset+self.HEADER_STRUCT_SIZE+length]
        return self.HEADER_STRUCT_SIZE + length

class DirectTCPSessionMessage(NMBSessionMessage):

    HEADER_STRUCT_FORMAT = '>I'
    HEADER_STRUCT_SIZE = struct.calcsize(HEADER_STRUCT_FORMAT)

    def decode(self, data, offset):
        data_len = len(data)

        if data_len < offset + self.HEADER_STRUCT_SIZE:
            # Not enough data for decoding
            return 0

        self.reset()
        length = struct.unpack(self.HEADER_STRUCT_FORMAT, data[offset:offset+self.HEADER_STRUCT_SIZE])[0]

        if length >> 24 != 0:
            raise NMBError("Invalid protocol header for Direct TCP session message")

        if data_len < offset + self.HEADER_STRUCT_SIZE + length:
            return 0

        self.data = data[offset+self.HEADER_STRUCT_SIZE:offset+self.HEADER_STRUCT_SIZE+length]
        return self.HEADER_STRUCT_SIZE + length