import array
import itertools
import socket
import struct

class ScratchError(Exception): pass
class ScratchConnectionError(ScratchError): pass	

class Scratch(object):

    prefix_len = 4
    broadcast_prefix_len = prefix_len + len('broadcast ')
    sensorupdate_prefix_len = prefix_len + len('sensor-update ')

    msg_types = set(['broadcast', 'sensor-update'])

    def __init__(self, host='localhost', port=42001):
        self.host = host
        self.port = port
        self.socket = None
        self.connected = False
        self.connect()

    def __repr__(self):
        return "Scratch(host=%r, port=%r)" % (self.host, self.port)

    def _pack(self, msg):
        """
        Packages msg according to Scratch message specification (encodes and 
        appends length prefix to msg). Credit to chalkmarrow from the 
        scratch.mit.edu forums for the prefix encoding code.
        """
        n = len(msg) 
        a = array.array('c')
        a.append(chr((n >> 24) & 0xFF))
        a.append(chr((n >> 16) & 0xFF))
        a.append(chr((n >>  8) & 0xFF))
        a.append(chr(n & 0xFF))
        return a.tostring() + msg 

    def _extract_len(self, prefix):
        """
        Extracts the length of a Scratch message from the given message prefix. 
        """
        return struct.unpack(">L", prefix)[0]

    def _get_type(self, s):
        """
        Converts a string from Scratch to its proper type in Python. Expects a
        string with its delimiting quotes in place. Returns either a string, 
        int or float. 
        """
        # TODO: what if the number is bigger than an int or float?
        if s.startswith('"') and s.endswith('"'):
            return s[1:-1]
        elif s.find('.') != -1: 
            return float(s) 
        else: 
            return int(s)

    def _escape(self, msg):
        """
        Escapes double quotes by adding another double quote as per the Scratch
        protocol. Expects a string without its delimiting quotes. Returns a new
        escaped string. 
        """
        escaped = ''	
        for c in msg: 
            escaped += c
            if c == '"':
                escaped += '"'
        return escaped

    def _unescape(self, msg):
        """
        Removes double quotes that were used to escape double quotes. Expects
        a string without its delimiting quotes, or a number. Returns a new
        unescaped string.
        """
        if isinstance(msg, (int, float, long)):
            return msg

        unescaped = ''
        i = 0
        while i < len(msg):
            unescaped += msg[i]
            if msg[i] == '"':
                i+=1
            i+=1
        return unescaped

    def _is_msg(self, msg):
        """
        Returns True if message is a proper Scratch message, else return False.
        """
        if not msg or len(msg) < self.prefix_len:
            return False
        length = self._extract_len(msg[:self.prefix_len])
        msg_type = msg[self.prefix_len:].split(' ', 1)[0]
        if length == len(msg[self.prefix_len:]) and msg_type in self.msg_types:
            return True
        return False

    def _parse_broadcast(self, msg):
        """
        Given a broacast message, returns the message that was broadcast.
        """
        # get message, remove surrounding quotes, and unescape
        return self._unescape(self._get_type(msg[self.broadcast_prefix_len:]))

    def _parse_sensorupdate(self, msg):
        """
        Given a sensor-update message, returns the sensors/variables that were
        updated as a dict that maps sensors/variables to their updated values.
        """
        update = msg[self.sensorupdate_prefix_len:]
        parsed = [] # each element is either a sensor (key) or a sensor value
        curr_seg = '' # current segment (i.e. key or value) being formed
        numq = 0 # number of double quotes in current segment
        for seg in update.split(' ')[:-1]: # last char in update is a space
            numq += seg.count('"')
            curr_seg += seg
            # even number of quotes means we've finished parsing a segment
            if numq % 2 == 0: 
                parsed.append(curr_seg)
                curr_seg = ''
                numq = 0
            else: # segment has a space inside, so add back it in
                curr_seg += ' '
        unescaped = [self._unescape(self._get_type(x)) for x in parsed]
        # combine into a dict using iterators (both elements in the list
        # inputted to izip have a reference to the same iterator). even 
        # elements are keys, odd are values
        return dict(itertools.izip(*[iter(unescaped)]*2)) 

    def _parse(self, msg):
        """
        Parses a Scratch message and returns a tuple with the first element
        as the message type, and the second element as the message payload. The 
        payload for a 'broadcast' message is a string, and the payload for a 
        'sensor-update' message is a dict whose keys are variables, and values
        are updated variable values. Returns None if msg is not a message.
        """
        if not self._is_msg(msg):
            return None
        msg_type = msg[self.prefix_len:].split(' ')[0]
        if msg_type == 'broadcast':
            return ('broadcast', self._parse_broadcast(msg))
        else:
            return ('sensor-update', self._parse_sensorupdate(msg))

    def _write(self, data):
        """
        Writes string data out to Scratch
        """
        total_sent = 0
        length = len(data)
        while total_sent < length:
            try:
                sent = self.socket.send(data[total_sent:])
            except socket.error as (err, msg):
                self.connected = False
                raise ScratchError("[Errno %d] %s" % (err, msg))
            if sent == 0:
                self.connected = False
                raise ScratchConnectionError("Connection broken")
            total_sent += sent

    def _send(self, data):
        """
        Sends a message to Scratch
        """
        self._write(self._pack(data))

    def _read(self, size):
        """
        Reads size number of bytes from Scratch and returns data as a string
        """
        data = ''
        while len(data) < size:
            try:
                chunk = self.socket.recv(size-len(data))
            except socket.error as (err, msg):
                self.connected = False
                raise ScratchError("[Errno %d] %s" % (err, msg))
            if chunk == '':
                self.connected = False
                raise ScratchConnectionError("Connection broken")
            data += chunk
        return data

    def _recv(self):
        """
        Receives and returns a message from Scratch
        """
        prefix = self._read(self.prefix_len)
        msg = self._read(self._extract_len(prefix))
        return prefix + msg

    def connect(self):
        """
        Connects to Scratch.
        """
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            self.socket.connect((self.host, self.port))
        except socket.error as (err, msg):
            self.connected = False
            raise ScratchError("[Errno %d] %s" % (err, msg))
        self.connected = True

    def disconnect(self):
        """
        Closes connection to Scratch
        """
        try: # connection may already be disconnected, so catch exceptions
            self.socket.shutdown(socket.SHUT_RDWR) # a proper disconnect
        except socket.error:
            pass
        self.socket.close()
        self.connected = False

    def sensorupdate(self, data):
        """
        Given a dict of sensors and values, updates those sensors with the 
        values in Scratch.
        """
        if not isinstance(data, dict):
            raise TypeError('Expected a dict')
        msg = 'sensor-update '
        for key in data.keys():
            msg += '"%s" "%s" ' % (self._escape(str(key)), 
                                    self._escape(str(data[key])))
        self._send(msg)

    def broadcast(self, msg):
        """
        Broadcasts msg to Scratch. msg can be a single message or an iterable 
        (list, tuple, set, generator, etc.) of messages.
        """
        if getattr(msg, '__iter__', False): # iterable
            for m in msg:
                self._send('broadcast "%s"' % self._escape(str(m)))
        else: # probably a string or number
            self._send('broadcast "%s"' % self._escape(str(msg)))

    def receive(self):
        """
        Receives broadcasts and sensor updates from Scratch. Returns a tuple
        with the first element as the message type and the second element 
        as the message payload. broadcast messages have a string as payload, 
        and the sensor-update messages have a dict as payload. Returns None if 
        message received could not be parsed. Raises exceptions on connection
        errors.
        """
        return self._parse(self._recv())