#Time to bring in my favorite friend, Twisted

import os, threading
import oommfq_core as qcore

from twisted.internet.protocol import Factory
from twisted.internet import reactor
from twisted.spread.banana import Banana, encode
from twisted.spread.jelly import jelly, unjelly

SendingSimsLock = threading.Lock()

def encodeTuplizedSim(tuplized):
    #I am *pretty sure* this is safe enough
    return (tuplized[0].encode("utf-8"),tuplized[1].encode("utf-8"), tuplized[2], tuplized[3])

##########################
# SERVER STATE CONSTANTS #
##########################

HANDSHAKE_NONE = -1
EXPECT_UID = 0
EXPECT_BATCHNAME = 1
EXPECT_FILENAME = 2
EXPECT_CHUNKNUM = 3
EXPECT_CHUNK = 4

class ServerListen(Banana):
    def __init__(self):
        Banana.__init__(self, isClient = False)
        self.needsKilling = False
        self.state = HANDSHAKE_NONE

        #Package data - these go live when a connection is made. "None" should be considered an
        #error state during file transfer
        self.username = None
        self.batchname = None
        self.batchpath = None
        self.filename = None
        self.fhandle = None
        self.chunk = None
        self.maxchunks = None

        self.simsToEnqueue=[]

    def connectionMade(self):
        #Banana has work to do...
        Banana.connectionMade(self)
        if not self.factory.serveManager.serverAlive:
            #Oops! Somehow someone made a connection when they weren't supposed to. Hang up on them.
            self.send("Sorry! This server is not currently friendly. This should never happen.")
            self.transport.loseConnection()
        else:
            self.send("OHAI")
            self.factory.sessions.add(self)
            self.state = EXPECT_UID

    def needsKilling(self):
        #Flag connection for destruction - unless it's in the middle of a file transfer.
        #Let that finish. The server is down, and this connection just needs to clean up.
        self.needsKilling = True
        if self.state < EXPECT_CHUNK:
            self.transport.loseConnection()

    def expressionReceived(self, sexp):
        data = unjelly(sexp)

        #First, check for netstat. Then, deal with all other processes.

        if data == "NETSTAT":
            #Send data on the running processes so a copy of OOMMFLink can visualize them.
            runSims = [encodeTuplizedSim(i.tuplize()) for i in self.factory.dataManager.runningSims]
            queueSims = [encodeTuplizedSim(i.tuplize()) for i in self.factory.dataManager.queuedSims]
            doneSims = [encodeTuplizedSim(i.tuplize()) for i in self.factory.dataManager.doneSims]
            errSims = [encodeTuplizedSim(i.tuplize()) for i in self.factory.dataManager.erroredSims]
            
            self.send(("NETSTAT", {"runningSims": runSims, "queuedSims": queueSims, 
                                   "doneSims":doneSims, "erroredSims":errSims}))

        elif self.state == EXPECT_UID:
            #Make sure we got a UID, and sent it.
            if len(data) == 2 and data[0] == "UID":
                self.username = data[1].decode("utf-8")
                self.state = EXPECT_BATCHNAME
                self.send("OK UID")
            else:
                self.send("ERR_EXPECT UID")

        elif self.state == EXPECT_BATCHNAME or (self.state == EXPECT_FILENAME and data[0] == "BATCH"):
            #In case we're resetting from a finished batch, clean house for safety's sake.
            self.filename = None
            self.chunk = None
            self.maxchunks = None

            if len(data) == 2 and data[0] == "BATCH":
                batchname = data[1].decode("utf-8")
                #Also, we should make dirs for it. Tab-split, and be sure to process directories VERY CAREFULLY on the way in...
                #Made directories go in ./remote/username/batchname
                madeDirs = self.factory.dataManager.prepareDirectoryForRemote([self.username] + batchname.split("\t"))
                if madeDirs:
                    #It worked!
                    self.send("OK BATCH")
                    self.state = EXPECT_FILENAME
                    self.batchname = batchname
                    self.batchpath = madeDirs
                else: #Uh-oh, makedirs whiffed. That path exists!
                    #Let the client sort things out and send us a new message. We'll drop you until then.
                    self.send("WARN_DATA EXISTS_BATCHPATH")
                    #self.factory.manager.killConnection()
                    self.transport.loseConnection()
            else:
                self.send("ERR_EXPECT BATCH")

        elif self.state == EXPECT_FILENAME:
            if len(data) == 2 and data[0] == "FNAME":
                self.filename = data[1].decode("utf-8")
                self.state = EXPECT_CHUNKNUM
                self.send("OK FNAME")
            elif data == "OK ALL START":
                self.sendAllSims()
            else:
                self.send("ERR_EXPECT FNAME")

        elif self.state == EXPECT_CHUNKNUM:
            #Here's where it gets good - how many chunks are we expecting?
            #Let's also write down that we're currently on chunk 0.
            if len(data) == 2 and data[0] == "CHUNKNUM":
                self.maxchunks = int(data[1])
                self.chunk = 0
                self.fhandle = open(self.batchpath + os.path.sep + self.filename, "wb") #Clobber mode.
                self.send("OK CHUNKNUM")
                self.state = EXPECT_CHUNK
            else:
                self.send("ERR_EXPECT CHUNKNUM")

        elif self.state == EXPECT_CHUNK:
            if len(data) == 3 and data[0] == "CHUNK":
                if not int(data[1]) == self.chunk:
                    #That's not what I wanted!
                    self.send("ERR_EXPECT CHUNK %s" % self.chunk)
                else:
                    #OK, start writing!
                    thischunk = data[2]
                    self.fhandle.write(thischunk)
                    #Now, determine whether or not to clean up.
                    self.chunk += 1 #One chunk is the zeroth chunk - might as well stop a fencepost error now.
                    if not self.chunk == self.maxchunks:
                        self.send("OK CHUNK")
                    else:
                        self.fhandle.close()
                        self.send("OK FILE")
                        #Snap back to filename-readiness to get the next part in the set.
                        self.chunk = None
                        self.maxchunks = None
                        if self.filename[-4:] == ".mif":
                            self.simsToEnqueue.append((qcore.localPath(self.batchpath,self.filename), self.username))
                        self.state = EXPECT_FILENAME

    def connectionLost(self, reason):
        self.sendAllSims()

    def sendAllSims(self):
        with SendingSimsLock:
            for sim in self.simsToEnqueue:
                self.factory.dataManager.enqueueSim(sim[0], sim[1])
            self.simsToEnqueue = []


    def send(self, msg):
        Banana.sendEncoded(self, jelly(msg))

class ServerListenFactory(Factory):
    protocol = ServerListen

    def __init__(self, dataManager, serveManager):
        #The most import part is that we have access to the data!
        self.dataManager = dataManager
        self.serveManager = serveManager
        self.sessions = set()

    def killall(self, gentle = True):
        #We want the server off. All existing connections can die.
        #Gentle mode (default) will not immediately tank connections
        for session in self.sessions:
            if gentle:
                session.needsKilling = True
            else:
                session.transport.loseConnection()


class _ServeManager(object):
    #This is the singleton object that manages starting and stopping the server
    def __init__(self):
        self.serverAlive = False
        self.server = None
        #This singleton object knows about another singleton object. It's cleaner this way.
        self.serverListenFactory = ServerListenFactory(qcore.OOMMFq, self)

        if int(qcore.OOMMFq.conf["autoServe"]) == 1:
            self.startServing()

        #Start the Twisted reactor in a lazy mode - use a wx timing loop to update it.
        reactor.startRunning(installSignalHandlers = 0)

    def startServing(self):
        self.serverAlive = True
        #Read the port from the factory's link to the OOMMFq data. Got that?
        port = int(self.serverListenFactory.dataManager.conf["port"])
        self.server = reactor.listenTCP(port, self.serverListenFactory)

    def stopServing(self):
        self.serverAlive = False
        if self.server:
            self.server.stopListening()
            self.serverListenFactory.killall() #Tell all connections to die *gently*. File transfers can finish.

    def netDoWork(self):
        reactor.runUntilCurrent()
        reactor.doIteration(0)
        
    def reactorDown(self):
        #I bet you want to quit the program.
        reactor.stop()

ServeManager = _ServeManager()
