#!/usr/bin/env python

import readline
import sys
import utils

# --------------------
#  Protocol utilities
# --------------------

class ProtocolError(Exception):
    '''Error caused by server not following protocol.'''
    def __init__(self, expr):
        self.expr = expr
    def __str__(self):
        return self.__class__.__name__ + "(" + repr(self.expr) + ")"

class BaseProtocol:
    def write(self, data):
        '''Override in subclass for writing data to port/stream/whatever.'''
        raise NotImplementedError()
    
    def read(self, bytecount):
        '''Override in subclass for reading data from port/stream/whatever.'''
        raise NotImplementedError()
    
    def readline(self):
        '''Read one line from port. Returns string without trailing \n.
        Skips over empty lines.
        Subclasses may override this with a more efficient implementation.
        '''
        bytes = [self.read(1)]
        while bytes[-1] == '\n':
            bytes = [self.read(1)]
        while bytes[-1] != '\n':
            bytes.append(self.read(1))
        return ''.join(bytes[:-1])
    
    def __iter__(self):
        while True:
            yield self.readline()
    
    def start_command(self, commandline):
        '''Start executing a command. Just writes the command line to
        port, appending newline if not already there.
        '''
        if not commandline.endswith('\n'):
            commandline += '\n'
        self.write(commandline)
    
    def get_partial_response(self):
        '''Read a single line and parse it to status, multiline, data.'''
        line = self.readline()
        try:
            status = int(line[:3])
            multiline = (line[3] == '-')
            data = line[4:]
        except (ValueError, IndexError):
            raise ProtocolError(line)
        
        return status, multiline, data
    
    def get_response(self):
        '''Wait for any response. Returns status, data. Data is a list of
        the response lines.
        '''
        data = []
        multiline = False
        for line in self:
            try:
                status = int(line[:3])
                multiline = (line[3] == '-')
                data.append(line[4:])
            except (ValueError, IndexError):
                if not multiline:
                    raise ProtocolError(line)
            
            if not multiline:
                break
        return status, data
    
    def get_final_response(self):
        '''Call get_response until the status code indicates it is a
        completion response (anything but 1xx).
        Returns status from final response and data from the all responses.
        '''
        alldata = []
        while True:
            status, data = self.get_response()
            alldata += data
            if not 100 <= status <= 199:
                return status, alldata
    
    def run_command(self, commandline):
        '''Run a command and return the status, data from get_final_response.'''
        self.start_command(commandline)
        return self.get_final_response()

class SerialProtocol(BaseProtocol):
    '''Serial port protocol.
    '''
    def __init__(self, port, baudrate, **kwargs):
        self.port = serial.Serial(port, baudrate, **kwargs)
        self.port.flushOutput()
        self.port.flushInput()
        
        self.port.setTimeout(0.5)
        try:
            self.port.read()
        except:
            pass
        try:
            self.run_command("")
        except:
            # Probably stuck in PUT or something
            self.port.write('\0' * 10000)
            self.run_command("")
        self.port.setTimeout(120.0)
    
    def write(self, bytes):
        self.port.write(bytes)
    
    def read(self, count):
        return self.port.read(count)

class InteractiveProtocol(SerialProtocol):
    '''Print every line as it is received by readline().'''
    def readline(self):
        data = SerialProtocol.readline(self)
        print data.rstrip('\n')
        return data

def is_ok(code):
    '''Return True if the code specifies a positive reply.'''
    return 200 <= code <= 299

# ------------------------------
#  Client-side command handlers
# ------------------------------

def handler_default(port, cmd, args):
    port.run_command(cmd + ' ' + ' '.join(args) + '\n')
    
def handler_get(port, cmd, args):
    if len(args) == 1:
        remote = args[0]
        local = args[0].split('/')[-1]
    elif len(args) == 2:
        remote, local = args
    else:
        print "Use: get remote [local]"
        return
    
    outfile = open(local, 'wb')
    port.start_command("get " + remote)
    status, multiline, data = port.get_partial_response()
    
    if status != 100:
        print "Unexpected GET status:", status
        return
    
    bytes = int(data)
    for block in utils.read_blocks(port, bytes):
        outfile.write(block)
        sys.stdout.write('.')
        sys.stdout.flush()
    sys.stdout.write('\n')
    
    status, data = port.get_final_response()
    if not is_ok(status):
        print "Error:", status
    return

def handler_put(port, cmd, args):
    if len(args) == 1:
        remote = args[0]
        local = args[0].split('/')[-1]
    elif len(args) == 2:
        local, remote = args
    else:
        print "Use: put local [remote]"
        return
    
    infile = open(local, 'r')
    infile.seek(0, 2)
    filesize = infile.tell()
    infile.seek(0)
    
    port.start_command("put " + remote + " " + str(filesize))
    
    for block in utils.read_blocks(infile, filesize, blocksize = 512):
        status, data = port.get_response()
        if status != 100:
            print "Unexpected PUT status:", status
            return
        port.write(block)
    
    status, data = port.get_final_response()
    if not is_ok(status):
        print "Error:", status
    return


special_handlers = {
    'put': handler_put,
    'get': handler_get,
}


# --------------
#  Main program
# --------------

class Completion:
    def __init__(self, port):
        self.commands = special_handlers.keys()
        
        # Request list of commands
        status, data = port.run_command("help")
        
        if is_ok(status):
            self.commands += data[0].strip().split()
        
        # Strip duplicates
        self.commands = list(set(self.commands))
    
    def complete(self, text, state):
        matches = [c for c in self.commands if c.startswith(c)]
        if len(matches) > state:
            return matches[state]
        else:
            return None
    
def main(port):
    completer = Completion(port)
    readline.set_completer(completer.complete)
    
    while True:
        line = raw_input("> ")
        parts = line.split(' ')
        cmd = parts[0]
        args = parts[1:]
        
        handler = special_handlers.get(cmd, handler_default)
        handler(port, cmd, args)

if __name__ == '__main__':
    import serial
    import sys

    if len(sys.argv) != 3:
        sys.stderr.write("Usage: " + sys.argv[0] + " /dev/ttyUSB0 115200\n")
        sys.exit(1)

    portname = sys.argv[1]
    baudrate = int(sys.argv[2])
    port = InteractiveProtocol(portname, baudrate)
    
    main(port)