import sys, os, time
from twisted.python import util
from twisted.internet import reactor, defer, task
from twisted.spread import pb

DEBUG = True

EXPIRE_INTERVAL = 5                    # Frequency (in seconds) to run
                                       # expiration checks.

class ConnectionRec:
    STATUS_CONNECTING = 0
    STATUS_CONNECTED = 1
    def __init__(self, hostport, expire):
        self.hostport = hostport
        self.expire = expire
        self.obj = None
        self.status = ConnectionRec.STATUS_CONNECTING
        self.callbacks = []
        
    def updateExpire(self, expireTime):
        self.expire = time.time() + expireTime
            

class ConnectionCache:
    """The connection cache maintains Perspective Broker connections,
    providing access to the root objects of remote
    machines. Connections are opened upon the first request, and kept
    open for subsequent requests until a timeout is reached."""

    
    def __init__(self, expireTime):
        self.__expireTime = expireTime
        self.__connections = {}
        self.__task = task.LoopingCall(self.expireConnections)
        self.__task.start(EXPIRE_INTERVAL)


    def connect(self, host, port):
        """Open a connection to the specified host and port to obtain
        the root object, or use the cached connection if one
        exists."""
        if (host, port) in self.__connections:
            connRec = self.__connections[host, port]
            if connRec.status == ConnectionRec.STATUS_CONNECTED:
                if DEBUG:
                    print "Using cached connection to", host, port
                # Already connected, just return the existing connection
                connRec.updateExpire(self.__expireTime)
                return defer.succeed(connRec.obj)
            elif connRec.status == ConnectionRec.STATUS_CONNECTING:
                # Currently attempting to connect. Defer to when the
                # connection is opened.
                if DEBUG:
                    print "Waiting for cached connection to", host, port
                dfd = defer.Deferred()
                connRec.updateExpire(self.__expireTime)
                connRec.callbacks.append(dfd)
                return dfd
            else:
                return defer.fail("BUG: invalid status for connection record")
        else:
            # Open a new connection
            if DEBUG:
                print "Opening connection to", host, port
            connRec = ConnectionRec((host,port),
                                    time.time() + self.__expireTime)
            factory = pb.PBClientFactory()
            reactor.connectTCP(host, port, factory)
            obj = factory.getRootObject()
            obj.addCallback(lambda x: self.gotConnection(connRec, x))
            obj.addErrback(lambda x: self.lostConnection(connRec, None))
            self.__connections[host, port] = connRec
            dfd = defer.Deferred()
            connRec.callbacks.append(dfd)
            return dfd

    def gotConnection(self, connRec, obj):
        """A connection has been established. Update state and call
        callbacks."""
        if DEBUG:
            print "Connection established to", connRec.hostport
        connRec.obj = obj
        connRec.status = ConnectionRec.STATUS_CONNECTED
        connRec.updateExpire(self.__expireTime)
        for x in connRec.callbacks:
            x.callback(obj)
        connRec.callbacks = []
        connRec.disconnectCB = lambda x : self.lostConnection(connRec, obj)
        obj.notifyOnDisconnect(connRec.disconnectCB)

    def lostConnection(self, connRec, obj):
        """A connection has been lost. Remove it from the cache."""
        if DEBUG:
            print "Connection lost from", connRec.hostport
        for cb in connRec.callbacks:
            cb.errback(pb.DeadReferenceError())
        del self.__connections[connRec.hostport]

    def expireConnections(self):
        curTime = time.time()
        for hostport, connRec in self.__connections.items():
            if connRec.expire < curTime:
                connRec.obj.dontNotifyOnDisconnect(connRec.disconnectCB)
                connRec.obj = None
                connRec.disconnectCB = None
                del self.__connections[hostport]