"""
Small standalone module implementing and supporting a simple KATCP
client specifically intended for use with ROACH.

This module should work with python 2.4, and since it's standalone
it's well suited to running on some fairly backwards platforms
(including ROACH with an ancient debian etch image)
"""
import socket

class ProgdevFailure(Exception):
    pass

class ListdevFailure(Exception):
    pass
    
class ListbofFailure(Exception):
    pass
    
class WriteFailure(Exception):
    pass

class ReadFailure(Exception):
    pass
    
class ConnectionFailure(Exception):
    pass
    
class BlockingClient(object):
    """
    A lightweight class that implements a simple blocking KATCP
    client directly over the socket interface.
    
    Currently this class only provides ROACH specific requests, however,
    there is no reason this can't be extended to include the full list
    of common (non ROACH) KATCP commands.
    
    Note the optional inspect_reply argument in many methods. This
    defaults to False, but if set to True, that method will blindly
    and immediately return the KATCP server's response as a string.
    KATCP is a human readable format, so this should be useful to 
    quickly identify bugs in this client.
    
    
    USAGE:
    
    This class should be very similar to corr.katcp_wrapper.FpgaClient;
    it essentially just offers a python interface to the KATCP requests
    described at https://casper.berkeley.edu/wiki/KATCP
    
    All one needs to do is import this file, instantiate a BlockingClient,
    and start making KATCP requests. For example:
    
        from roach_katcp import *
        
        roach = BlockingClient('some_ip_address')
        
        roach.progdev('some_file')
        
        data = roach.read('some_register',4)
    
    TODO:
    
    Implement wordwrite, wordread, and tap_start
    
    """
    
    def __init__(self, host_ip = 'localhost', port = 7147):
        """
        Setup a client socket on the katcp port
        """
        # max number of bytes to recieve for replies to katcp
        # requests that don't involve FPGA storage, such as
        # listdev and progdev
        self.buffsize = 4096
        
        self.host_ip = host_ip
        self.port = port
        
        try:
            self.katcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.katcp_socket.connect((host_ip, port))
        
            # tcpborhserver sends a small banner message on connecting
            self.katcp_socket.recv(self.buffsize)
        except:
            msg = 'Could not connect to %s:%s\n'%s(self.host_ip,self.port)
            raise ConnectionFailure(msg)
        
    def test(self):
        """
        Check the connection to the server with a "?help" request.
        
        Returns True if conection is good, else False
        """
        
        try:
            self.katcp_socket.send('?help\n')
            reply = self.katcp_socket.recv(self.buffsize)
            
        except:
            return False
            
        if not('!help ok' in reply):
            return False
            
        return True
    
    def reconnect(self):
        """
        Close and re-establish the katcp connection. Useful for purging
        any unwanted data queued up by the OS for recv.
        
        """
        try:
            self.katcp_socket.close()
            self.katcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.katcp_socket.connect((self.host_ip, self.port))
            # tcpborhserver sends a small banner message on connecting
            self.katcp_socket.recv(self.buffsize)
            self.test()
        except:
            msg = 'Could not reconnect to %s:%s\n'%s(self.host_ip,self.port)
            raise ConnectionFailure(msg)
        
    def send(self,msg,buffsize = 4096, flush = True):
        """
        Send a string directly to the KATCP server and blindly
        return the reply. Newline character is automatically appended
        to the message string.
        
        Default buffer size is 4096 bytes. This should be much more
        than sufficient for anything except a large read from a bram.
        
        
        Notes:
        
        If flush = True:
        
        Socket is closed and then re-established after recv. This
        protects subsequent recv's from other class methods from 
        unwittingly fetching leftover data and failing.        
        
        If flush = False:
        
        Note that using a buffer that is too small will not cause an
        error in this function, but can cause an error in a subsequent
        KATCP request: send() will return as much of the message as
        will fit in buffsize bytes and exit normally, but will leave
        behind the rest of the message for future recv calls to absorb.
        This can cause future calls to BlockingClient methods to fail
        when they pull leftover response data from the socket and parse
        it incorrectly.
        """
        self.katcp_socket.send(msg+'\n')
        reply = self.katcp_socket.recv(buffsize)
        
        if flush:
            self.reconnect()  
        
    def progdev(self, filename, inspect_reply = False):
        """
        Attempt to load a boffile. Return the process id if
        successful.
        """
        
        # dispatch katcp request and wait for reply
        msg = '?progdev %s\n' % (filename)
        self.katcp_socket.send(msg)
        reply = self.katcp_socket.recv(self.buffsize)
        
        if inspect_reply:
            return reply
            
        # progdev is strange--tcpborhpserver will reply "!progdev ok"
        # even if we try to progdev a nonexistent file. I actually
        # have never managed to generate a "!progdev fail" so I have
        # no idea what if anything "!progdev ok" means. Nonetheless, we
        # check for it here. Luckily, when things do go wrong, we seem
        # to get a "#log warn".
        
        if 'ok' in reply:
            progdev_ok = True
        else:
            progdev_ok = False
            
        if 'warn' in reply:
            no_warning = False
        else:
            no_warning = True
            
        ok = progdev_ok and no_warning
        
        if not ok:
            raise ProgdevFailure('Could not load bof %s'%(filename))

        pid = int(reply.split()[2])
            
        return pid
            
    def listdev(self, inspect_reply = False):
        """
        Return a list of names as strings of available FPGA registers
        """
        
        # dispatch kactp request and wait for reply
        msg = '?listdev\n'
        self.katcp_socket.send(msg)
        reply = self.katcp_socket.recv(self.buffsize)
        
        if inspect_reply:
            return reply
        
        if 'fail' in reply:
            raise ListdevFailure()
            
        # split the reply into a list of strings representing each line.
        # The last line contains a "!listdev ok" and is discarded.
        lines = reply.split('\n')[0:-2]
        
        out = []
        for line in lines:
            # we're interested in the second word on every line,
            # the first word on the line should be "#listdev"
            out.append(line.split()[1])
            
        return out
        
    def listbof(self, inspect_reply = False):
        """
        Return a list of names as strings of .bof files in /boffiles
        """
        
        # everything is the same as for listdev

        msg = '?listbof\n'
        self.katcp_socket.send(msg)
        reply = self.katcp_socket.recv(self.buffsize)
        
        if inspect_reply:
            return reply
        
        # admittedly, I don't know what the error message for
        # listbof looks like...
        if 'fail' in reply:
            raise ListbofFailure()
            
        lines = reply.split('\n')[0:-2]
        
        out = []
        for line in lines:
            out.append(line.split()[1])
            
        return out
        
    def status(self, inspect_reply = False):
        """
        Return True if is FPGA is programmed
        """
        msg = '?status\n'
        self.katcp_socket.send(msg)
        reply = self.katcp_socket.recv(self.buffsize)
        
        if inspect_reply:
            return reply
        
        if 'ok' in reply:
            return True
        return False
        
    def write(self, register, data, offset = 0, inspect_reply = False):
        """
        Attempts to write a register. Does not return anything unless
        inspect_reply = True
        """
        
        msg = '?write %s %s %s\n' % (register,offset,data)
        self.katcp_socket.send(msg)
        reply = self.katcp_socket.recv(self.buffsize)
        
        if inspect_reply:
            return reply
        
        if 'fail' in reply:
            msg = ', bad register name?'
            if 'program' in reply:
                msg = ', FPGA not programmed?'
            raise WriteFailure('Failed to write %s%s' %(register,msg))
            
    def read(self, register, size, offset = 0, inspect_reply = False):
        """
        Return data at register as a binary big endian encoded string.
        """
        
        msg = '?read %s %s %s\n' % (register,offset,size)
        self.katcp_socket.send(msg)
        reply = self.katcp_socket.recv(size+128)
        
        if inspect_reply:
            return reply
        
        if 'fail' in reply:
            msg = ', bad register name?'
            if 'program' in reply:
                msg = ', FPGA not programmed?'
            raise ReadFailure('Failed to read %s%s' %(register,msg))
            
        # reply should be formatted as "!read ok <datastring>"
        return reply.split()[2]
        
        
    def __del__(self):
        
        self.katcp_socket.close()