#!/usr/bin/env/python
# Tells Linux systems how to run this file.

"""
Toomi
A socket-based server/client system.

Thanks to the tutorial at <http://www.amk.ca/python/howto/sockets/>
and to Enrico Kochon of the Pygame discussion list.
"""

__author__ = "Kris Schnee"
__version__ = "2010.5.23"
__license__ = "GPLv3 or MIT"

##### Imported Modules #####

## Standard

import socket
import select
import time

## Third-party


## Mine


##### Constants #####

HOST = "localhost"
PORT = 7777

##### Classes #####

class SocketBase:
    """Contains shared send/receive functions for client & server classes.

    Note that both functions assume a 6-byte header string giving the length
    of the data to follow. So if you use one function, you must use the other,
    always, or you risk misinterpretation of your data."""

    def __init__(self):
        self.socket = None ## Will give an error unless overridden.

    def Send(self,data,socket=None):
        socket = socket or self.socket ## Default
        if not data:
            return
        size = len(data)
        if size > 999999:
            raise SystemError("Can't send more than 999999 bytes at once.")
        size_str = str(size).zfill(6) ## Eg. "000042"
        output = size_str + data
##        print "Sending "+str(size+6)+" bytes."

        """Send as many bytes as the connection allows, then keep trying to
        send from where we left off, till we reach the end."""
        total_sent = 0
        while total_sent < size:
            sent = socket.send(output[total_sent:])
            if not sent: ## Nothing got through?
                raise SystemError("Connection seems to have broken.")
            total_sent += sent

    def Receive(self,socket=None):
        socket = socket or self.socket ## Default
        size = int(socket.recv(6))
##        print "Receiving "+str(size)+" bytes."
        total_received = 0
        data = ""
        while total_received < size:
            new_data = socket.recv(size - total_received)
            if not new_data:
                raise SystemError("Expected to receive "+str(size)+" bytes, but only got "+str(total_received)+". Broken connection?")
            data += new_data
            total_received += len(data)
        return data


class Server( SocketBase ):
    """A socket-based Net server for multiple clients.

    This class doesn't send anything to clients. It establishes connections
    from any client that contacts it. It also tracks a little info about
    each client. Use a subclass of this to do things with the input, such as
    relaying anything one client sends to all other clients. See demo class."""
    
    def __init__(self,internet_host=True,port=PORT):
        SocketBase.__init__(self)
        self.clients = {}

        self.listen_socket = socket.socket( socket.AF_INET, socket.SOCK_STREAM )
        self.listen_socket.setblocking(False)
        if internet_host:
            host = socket.gethostbyname(socket.gethostname())
        else:
            host = "localhost"
        self.listen_socket.bind( (host, port) )
        self.listen_socket.listen(4)
        print "Server started. My IP address and port: "+str(host)+":"+str(port)
        f = open("server_host_info.txt","w")
        f.write("Server started. My IP address and port: "+str(host)+":"+str(port))
        f.close()

    def ShutDownServer(self,end_program=True):
        self.listen_socket.close()
        for c in self.clients.values():
            c["socket"].close()
        if end_program:
            exit()

    def HandleInput(self,address,data):
        """Override this."""
        result = "Input received: "+data
        if data == "quit":
            self.HandleQuitRequest(address)
        else:
            self.Send(result,self.clients[address]["socket"])

    def CheckForConnections(self):
        """Check for connection attempts."""
        rejoiners = [] ## Known addresses, reconnecting for some reason.
        newcomers = []
        
        a = [self.listen_socket]
        b = []
        c = []
        readable, writeable, erroneous = select.select(a,b,c,.001)
        if readable:
##            print "Listen socket seems to be readable. Checking now."
            socket, address = self.listen_socket.accept()
            print "Got a connection from "+str(address)
            socket.setblocking(False)
            
            if address in self.clients:
                print "Rejoining client: "+str(address)
                rejoiners.append((socket,address))
            else:
                print "New client: "+str(address)
                newcomers.append((socket,address))
        return rejoiners, newcomers

    def HandleNewConnections(self,rejoiners=[],newcomers=[]):
        """React to newly connected clients. Override this."""

        for socket, address in rejoiners:
            self.clients[address]["socket"].close() ## Close old socket
            self.clients[address]["socket"] = socket ## Replace socket
            self.clients[address]["last_input"] = int(time.time)

        for socket, address in newcomers:
            self.clients[address] = {"last_input":int(time.time()),
                                     "name":"stranger",
                                     "socket":socket,
                                     }

    def HandleQuitRequest(self,client):
        """Client requested to quit. Clean up the connection.

        In a more complex program, override this function to do
        further cleanup like removing a character from a game world."""

        print "Got a quit request from: "+str(client)
        
        if not client in self.clients:
            raise SystemError("Somehow I was sent a nonexistent client.")
        self.clients[client]["socket"].close()
        del self.clients[client] ## Stop listening to this one.

    def CheckForInput(self):
        """Get input from clients already connected."""
        inputs = {}
        a = [c["socket"] for c in self.clients.values()]
        b = []
        c = []
        if a:
            readable, writeable, erroneous = select.select(a,b,c,.001)
            for client in self.clients:
                s = self.clients[client]["socket"]
                if s in readable:
                    ## Socket is ready to send to me. Receive now.
                    data = self.Receive(s)
                    inputs[client] = data ## Record "client X says Y".
        return inputs

    def DemoLoop(self):
        while True:
            rejoiners, newcomers = self.CheckForConnections()
            if rejoiners or newcomers:
                self.HandleNewConnections(rejoiners,newcomers)
                
            inputs = self.CheckForInput()
            for address in inputs:
                if inputs[address]:
                    ## Act based on the input, maybe replying.
                    self.HandleInput(address,inputs[address])
            time.sleep(.1)




class Client( SocketBase ):
    def __init__(self):
        SocketBase.__init__(self)
        self.socket = None

    def Connect(self,host=HOST,port=PORT):
        self.socket = socket.socket( socket.AF_INET, socket.SOCK_STREAM )
        self.socket.connect( (HOST, PORT) )
        self.socket.setblocking(0)

    def Disconnect(self):
        self.socket.close()

    def ListenToServer(self):
        """Get input."""
        readable, writeable, erroneous = select.select([self.socket],[],[])
        if readable:
            data = self.Receive(self.socket)
            return data

    def DemoLoop(self,host=HOST,port=PORT):
        print "Starting client program."
        self.Connect(host,port)
        while True:
            output = raw_input("Send what? > ")

            if output.lower() in ["q","quit","exit"]:
                print "Quitting."
                self.Send("quit")
                import time
                time.sleep(2)
                return
            
            self.Send(output)
            reply = self.ListenToServer()
            print reply



class ServerDemo( Server ):
    def __init__(self,use_internet=True,port=PORT):
        Server.__init__(self,use_internet,port)
        self.weather = "sunny"

    def HandleInput(self,address,data):
        """Demonstrate reacting to user input, eg. changing "weather"."""
        result = "Input received: "+data

        if data.lower() == "shutdown":
            self.ShutDownServer()
            return

        elif data.lower() == "quit":
            self.HandleQuitRequest(address)
            return

        name = self.clients[address]["name"]
        result = ""

        words = data.split(" ")
        if words:
            if words[0] == "weather":
                self.weather = words[1]

            elif words[0] == "name":
                self.clients[address]["name"] = " ".join(words[1:])

            elif words[0] == "say":
                result = name+" says: "+data+"\n\n"

        result += "Hello, "+name+". The weather is currently "+self.weather+".\nMy currently connected clients are: "
        result += ", ".join([c["name"] for c in self.clients.values()])
        result += "\nThe last thing you sent was "+str(len(data))+" bytes long."
        result += "\nYou can type 'quit', 'weather [something]', 'name [something]' or 'say [something]' to do stuff."

        self.Send(result,self.clients[address]["socket"])



##### Other Functions #####

def AskForIP():
    host = raw_input("Host IP address? (leave blank for localhost) > ")
    if not host:
        host = HOST
    port = raw_input("Host port? (leave blank for default) > ")
    if not port:
        port = PORT
    return host, port

def AskUseInternet():
    x = raw_input("Start an Internet server (y/n)? No/blank = local server only. > ")
    if x in "yY":
        return True
    return False

def GetMyIP():
    return socket.gethostbyname(socket.gethostname())

##### Autorun #####

if __name__ == "__main__":
    ## Runs if this module is the main module (not imported).
    use_internet = AskUseInternet()
    s = ServerDemo(use_internet)
    s.DemoLoop()