from optparse import OptionParser

import hl7
import os.path
import six
import socket
import sys

SB = b'\x0b'  # <SB>, vertical tab
EB = b'\x1c'  # <EB>, file separator
CR = b'\x0d'  # <CR>, \r

FF = b'\x0c'  # <FF>, new page form feed


class MLLPException(Exception):

class MLLPClient(object):
    A basic, blocking, HL7 MLLP client based upon :py:mod:`socket`.

    MLLPClient implements two methods for sending data to the server.

    * :py:meth:`MLLPClient.send` for raw data that already is wrapped in the
      appropriate MLLP container (e.g. *<SB>message<EB><CR>*).
    * :py:meth:`MLLPClient.send_message` will wrap the message in the MLLP

    Can be used by the ``with`` statement to ensure :py:meth:`MLLPClient.close`
    is called::

        with MLLPClient(host, port) as client:

    MLLPClient takes an optional ``encoding`` parameter, defaults to UTF-8,
    for encoding unicode messages [#]_.

    .. [#]
    def __init__(self, host, port, encoding='utf-8'):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.connect((host, port))
        self.encoding = encoding

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, trackeback):

    def close(self):
        """Release the socket connection"""

    def send_message(self, message):
        """Wraps a byte string, unicode string, or :py:class:`hl7.Message`
        in a MLLP container and send the message to the server

        If message is a byte string, we assume it is already encoded properly.
        If message is unicode or  :py:class:`hl7.Message`, it will be encoded
        according to  :py:attr:`hl7.client.MLLPClient.encoding`

        if isinstance(message, six.binary_type):
            # Assume we have the correct encoding
            binary = message
            # Encode the unicode message into a bytestring
            if isinstance(message, hl7.Message):
                message = six.text_type(message)
            binary = message.encode(self.encoding)

        # wrap in MLLP message container
        data = SB + binary + EB + CR
        return self.send(data)

    def send(self, data):
        """Low-level, direct access to the socket.send (data must be already
        wrapped in an MLLP container).  Blocks until the server returns.
        # upload the data
        # wait for the ACK/NACK
        return self.socket.recv(RECV_BUFFER)

# wrappers to make testing easier
def stdout(content):
    # In Python 3, can't write bytes via sys.stdout.write
    if six.PY3 and isinstance(content, six.binary_type):
        out = sys.stdout.buffer
        newline = b'\n'
        out = sys.stdout
        newline = '\n'

    out.write(content + newline)

def stdin():
    return sys.stdin

def stderr():
    return sys.stderr

def read_stream(stream):
    """Buffer the stream and yield individual, stripped messages"""
    _buffer = b''

    while True:
        data =
        if data == b'':
        # usually should be broken up by EB, but I have seen FF separating
        # messages
        messages = (_buffer + data).split(EB if FF not in data else FF)

        # whatever is in the last chunk is an uncompleted message, so put back
        # into the buffer
        _buffer = messages.pop(-1)

        for m in messages:
            yield m.strip(SB + CR)

    if len(_buffer.strip()) > 0:
        raise MLLPException('buffer not terminated: %s' % _buffer)

def read_loose(stream):
    """Turn a HL7-like blob of text into a real HL7 messages"""
    # look for the START_BLOCK to delineate messages
    START_BLOCK = b'MSH|^~\&|'

    # load all the data
    data =

    # take out all the typical MLLP separators. In Python 3, iterating
    # through a bytestring returns ints, so we need to filter out the int
    # versions of the separators, then convert back from a list of ints to
    # a bytestring (In Py3, we could just call bytes([ints]))
    separators = [six.byte2int(bs) for bs in [EB, FF, SB]]
    data = b''.join([six.int2byte(c) for c in six.iterbytes(data) if c not in separators])

    # Windows & Unix new lines to segment separators
    data = data.replace(b'\r\n', b'\r').replace(b'\n', b'\r')

    for m in data.split(START_BLOCK):
        if not m:
            # the first element will not have any data from the split

        # strip any trailing whitespace
        m = m.strip(CR + b'\n ')

        # re-insert the START_BLOCK, which was removed via the split
        yield START_BLOCK + m

def mllp_send():
    """Command line tool to send messages to an MLLP server"""
    # set up the command line options
    script_name = os.path.basename(sys.argv[0])
    parser = OptionParser(usage=script_name + ' [options] <server>')
        action='store_true', dest='version', default=False,
        help='print current version and exit'
        '-p', '--port',
        action='store', type='int', dest='port', default=6661,
        help='port to connect to'
        '-f', '--file', dest='filename',
        help='read from FILE instead of stdin', metavar='FILE'
        '-q', '--quiet',
        action='store_true', dest='verbose', default=True,
        help='do not print status messages to stdout'
        action='store_true', dest='loose', default=False,
            'allow file to be a HL7-like object (\\r\\n instead '
            'of \\r). Requires that messages start with '
            '"MSH|^~\\&|". Requires --file option (no stdin)'

    (options, args) = parser.parse_args()

    if options.version:
        import hl7

    if len(args) == 1:
        host = args[0]
        # server not present
        stderr().write('server required\n')

    if options.filename is not None:
        # Previously set stream to the open() handle, but then we did not
        # close the open file handle.  This new approach consumes the entire
        # file into memory before starting to process, which is not required
        # or ideal, since we can handle a stream
        with open(options.filename, 'rb') as f:
            stream = six.BytesIO(
        if options.loose:
            stderr().write('--loose requires --file\n')
        stream = stdin()

    with MLLPClient(host, options.port) as client:
        message_stream = (
            if not options.loose
            else read_loose(stream)

        for message in message_stream:
            result = client.send_message(message)
            if options.verbose:

if __name__ == '__main__':