from SimpleXMLRPCServer import SimpleXMLRPCServer
from SimpleXMLRPCServer import SimpleXMLRPCRequestHandler
import xmlrpclib
import pickle
try:
    import pyfits
except ImportError:
    import astropy.io.fits as pyfits
# improves performance a few percent
pyfits.ENABLE_RECORD_VALUED_KEYWORD_CARDS = False

import numpy
import cStringIO
def ser_array(x):
    if x is None:
        return x
    f = cStringIO.StringIO()
    numpy.lib.format.write_array(f, x)
    return xmlrpclib.Binary(f.getvalue())

def deser_array(s):
    if s == None:
        return None
    f = cStringIO.StringIO(s.data)
    return numpy.lib.format.read_array(f)


# Restrict to a particular path.
class RequestHandler(SimpleXMLRPCRequestHandler):
        rpc_paths = ('/RPC2',)

class Server:
    def __init__(self, port=8000):
        # Create server
        self.server = SimpleXMLRPCServer(("localhost", port),
                                         requestHandler=RequestHandler,
                                         allow_none=True,
                                         logRequests=False)
        self.server.register_introspection_functions()
        self.server.register_multicall_functions()
        self.server.register_function(self.quit)

        self.server.register_function(self.pyfits_open, 'pyfits_open')
        self.server.register_function(self.pyfits_close, 'pyfits_close')

        self.server.register_function(self.pyfits_getheader, 'pyfits_getheader')
        self.server.register_function(self.pyfits_getdata, 'pyfits_getdata')
        self.server.register_function(self.pyfits_getval, 'pyfits_getval')
        self.server.register_function(self.pyfits_custom, 'pyfits_custom')
        self.server.register_function(self.len, 'len')

        self.files = dict()

        self._run = True

    def run(self):
        while self._run:
            self.server.handle_request()

    def quit(self):
        self._run = False
        return 0

    def pyfits_open(self, fn):
        # while memory mapping would be good it is handled badly by pyfits
        # see tickets 175 and 177
        if fn not in self.files:
            # workaround pyfits <= 2.1.1 unicode issue (CNSHD750437)
            try:
                self.files[fn] = pyfits.open(fn, memmap=False)
            except AttributeError:
                self.files[fn] = pyfits.open(str(fn), memmap=False)

    def pyfits_close(self, fn):
        if fn in self.files:
            del self.files[fn]

    def pyfits_getheader(self, fn, idx):
        self.pyfits_open(fn)
        f = self.files[fn]
        if isinstance(idx, int):
            return dict(f[idx].header)
        if isinstance(idx, list):
            vals = []
            for i in idx:
                vals.append(dict(f[i].header))
            return vals
        if idx is None:
            return [dict(x.header) for x in f]

    def pyfits_getdata(self, fn, idx):
        self.pyfits_open(fn)
        f = self.files[fn]
        if isinstance(idx, int):
            return ser_array(f[idx].data)
        if isinstance(idx, list):
            vals = []
            for i in idx:
                vals.append(ser_array(f[i].data))
            return vals
        if idx is None:
            return [ser_array(x.data) for x in f]
        return ser_array(f[idx].data)

    def pyfits_getval(self, fn, keyword, idx, skip_first=True):
        self.pyfits_open(fn)
        f = self.files[fn]
        if isinstance(idx, int):
            return f[idx].header.get(keyword)
        if isinstance(idx, list):
            vals = []
            for i in idx:
                vals.append(f[i].header.get(keyword))
            return ser_array(numpy.array(vals))
        if idx is None:
            if skip_first:
                return ser_array(numpy.array([x.header.get(keyword) for x in f[1:]]))
            else:
                return ser_array(numpy.array([x.header.get(keyword) for x in f]))

    def pyfits_custom(self, fn, pfun, args=None):
        self.pyfits_open(fn)
        f = self.files[fn]
        fun = pickle.loads(pfun)
        if args:
            res = fun(f, *args)
        else:
            res = fun(f)
        if isinstance(res, numpy.ndarray):
            return ser_array(res)
        else:
            return res

    def len(self, fn):
        return self.pyfits_custom(fn, len)

if __name__ == "__main__":
    s = Server()
    s.run()
