#!/usr/bin/python
import logging
import struct
import bson
from six import BytesIO, b

class BSONClient(object):
    def __init__(self, sock):
        self.sock = sock
        self.log = logging.getLogger(self.__class__.__name__)

    @staticmethod
    def _bintoint(data):
        return struct.unpack("<i", data)[0]

    def send_packet(self, obj):
        """
            Atomically send a BSON message.
        """
        data = bson.dumps(obj)
        self.sock.sendall(data)

    def recv_packet(self):
        """
            Atomic read of a BSON message.
            This function either returns a dict, None, or raises a socket error.
            If the return value is None, it means the socket is closed by the other side.
        """
        sock_buf = self.recvbytes(4)
        if sock_buf is None:
            return None

        message_length = _bintoint(sock_buf.getvalue())
        self.log.debug("Receving Message (size=%d)", message_length)
        sock_buf = self.recvbytes(message_length - 4, sock_buf)
        if sock_buf is None:
            return None

        retval = bson.loads(sock_buf.getvalue())
        return retval


    def recvbytes(self, bytes_needed, sock_buf = None):
        """
            Atomic read of bytes_needed bytes.
            This function either returns exactly the nmber of bytes requested in a
            StringIO buffer, None, or raises a socket error.
            If the return value is None, it means the socket is closed by the other side.
        """
        if sock_buf is None:
           sock_buf = BytesIO()

        bytes_count = 0
        while bytes_count < bytes_needed:
            chunk = self.sock.recv(min(bytes_needed - bytes_count, 32768))
            part_count = len(chunk)
            self.log.debug("Received Data (size=%d)", part_count)

            if type(chunk) == str:
                chunk = b(chunk)

            if part_count < 1:
                return None

            bytes_count += part_count
            sock_buf.write(chunk)

        return sock_buf

    def __call__(self, client, address):
        while True:
            obj = self.recv_packet(client)
            self.log.debug("received object: %s", repr(obj))
            if "cmd" in obj:
                if hasattr(self,"do_"+obj["cmd"]):
                    retval = getattr(self,"do_"+obj["cmd"])(obj)
                    self.log.debug("sending object: %s", repr(obj))
                    if retval is not None:
                        self.send_packet(client,retval)
            else:
                self.send_packet(client,{"error": "No Command"})


if __name__ == '__main__':
    import sys
    import logging
    import socket
    import argparse

    sys.path.append("./lib")

    parser = argparse.ArgumentParser(description='Send a vmchannel command.')

    commands = parser.add_mutually_exclusive_group(required=True)
    commands.add_argument('-m', '--monitor',
                    dest='monitor', type=int, action='store',
                    help='change monitor')
    parser.add_argument('-s', '--socket',
                    dest='socket', type=str, action='store', default="/var/run/libvirt/vmchannel.sock",
                    help='Socket path')
    parser.add_argument('-v', '--verbose',
                    dest='verbose', action='count', default=0,
                    help='Multiple -v options increase the verbosity.  The maximum is 3.')

    args = parser.parse_args()

    root = logging.getLogger()
    root.setLevel(logging.DEBUG)
    ch = logging.StreamHandler()
    ch.setLevel(50-10*args.verbose)
    ch.setFormatter(logging.Formatter('%(asctime)s - %(name)s.%(funcName)s - %(levelname)s - %(message)s'))
    root.addHandler(ch)

    socket_client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    socket_client.connect(args.socket)

    class MonitorControl(object):
        def switch_monitor(self,monitor_number):
            return self.send_packet({"cmd": "switch_monitor", "monitor": monitor_number})

    class RunClient(MonitorControl,BSONClient):
        pass

    app = RunClient(socket_client)
    if 'monitor' in args:
        app.switch_monitor(args.monitor)