#!/usr/bin/env python
# Chat Program - chat.py
import socket, traceback, os, sys, select

class stateclass:
    stdmask = select.POLLERR | select.POLLHUP | select.POLLNVAL
    
    def __init__(self, mastersock):
        # Creates objects for stateclass
        # Creates a poll object
        self.p = select.poll()
        self.mastersock = mastersock
        self.watchread(mastersock)
        self.readbuffers = {}
        self.writebuffers = {}
        self.sockets = {mastersock.fileno(): mastersock}
    
    def fd2socket(self, fd):
        # Return a socket, given a file descriptor
        return self.sockets[fd]
    
    # The following 'watch' functions tell the operating system
    # what it should be expecting for the socket: reading, writing,
    # both, or none. The function names may be self-explanatory.
    def watchread(self, fd):
        self.p.register(fd, select.POLLIN | self.stdmask)
    def watchwrite(self, fd):
        self.p.register(fd, select.POLLOUT | self.stdmask)
    def watchboth(self, fd):
        self.p.register(fd, select.POLLIN | select.POLLOUT | self.stdmask)
    def dontwatch(self, fd):
        self.p.unregister(fd)
    
    # Called when a request for a new connection arrives
    def newconn(self, sock):
        fd = sock.fileno()
        self.watchboth(fd)
        self.readbuffers[fd] = ""
        self.sockets[fd] = sock
    
    # Called when new data arrives
    def readevent(self, fd):
        try:
            # Read in the data and throw it on the buffer!
            self.readbuffers[fd] += self.fd2socket(fd).recv(4096)
            print self.readbuffers[fd]
        except:
            self.closeout(fd) # Kill if it don't work!
        # Tells the system to watch for both reading and writing
        self.watchboth(fd)
            
    # Called when data is ready to be written
    def writeevent(self, fd):
        if not len(self.readbuffers[fd]):
            # If no data, removes from the list to write
            self.watchread(fd)
            return
        try:
            # self.fd2socket gets the socket information to send to
            # self.writebuffers[fd] is the data for 'send'
            byteswritten = self.fd2socket(fd).send(self.readbuffers[fd])
        except:
            self.closeout(fd)
        
        # Clears the buffer of the sent text
        self.readbuffers[fd] = self.readbuffers[fd][byteswritten:]
        
        # If the buffer is empty, limits socket to reading
        if not len(self.writebuffers[fd]):
            self.watchread(fd)
        
    # If anything screws up, it closes the socket
    def errorevent(self, fd):
        self.closeout(fd)
        
    # Closes the connections
    def closeout(self, fd):
        # Stops system for looking for read and writes
        self.dontwatch(fd)
        # Tries hard to kill that socket!
        # fd2socket(fd) specifies which socket
        try:
            self.fd2socket(fd).close()
        except:
            pass
        
        # Wipes the socket and the buffers
        del self.writebuffers[fd]
        del self.sockets[fd]
    
## END OF CONFUSING HEAP OF NETWORKING, BEGIN FUN CODE ##

# Starts up a server and waits for clients
def server(s):
    host = ''  # Binds to all, anyone can connect
    s.bind((host, port))  # Binds the server to a host add a port number
    s.listen(1)  # Tells the server to listen for connections
    s.setblocking(0)
    state = stateclass(s) # Creates class object
    print "Waiting for clients..."

    # Main function: Waits for socket activity, when it
    # recieves activity, it determines what the activity
    # is (reading, writing, new client), and then sends
    # it to the appropriate function.
    while 1:
        # Waits for some kind of socket activity
        result = state.p.poll()
        # Recieved activity:
        for fd, event in result:
            # If event is a new client connection (mastersock)
            if fd == state.mastersock.fileno() and event == select. POLLIN:
                try:
                    newsock, addr = state.fd2socket(fd).accept()
                    newsock.setblocking(0)
                    print "Recieved connection from %s, sir." % str(newsock.getpeername())
                    # Processes the connection in the newconn function
                    state.newconn(newsock)
                except:
                    pass
            # Sends the event to proper function
            elif event == select.POLLIN:
                state.readevent(fd)
            elif event == select.POLLOUT:
                state.writeevent(fd)
            else:
                state.errorevent(fd)

# Attempts to connect to a server
def client(s):
    host = raw_input('What server would you like to connect to?: ')
    s.connect((host, port))
    print "Connected to %s, sir" % str(s.getpeername()),
    print "Press Ctrl-C to exit, buddy."
    while 1:
        data = raw_input(name + ' says?: ')
        final = name + " said: %s" % data
        s.sendall(final)

port = 50000
# Create socket s(    IPv4      ,        TCP        )
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
# Setting Socket Options( level , optname , value)
# SO_REUSEADDR: Releases port immediately after server stops
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

## MAIN CODE ##
name = raw_input('What is your name, friend?: ')
while(1):
    question = raw_input('Would you like to (c)onnect to someone, or (w)ait for a connection?: ')
    if question == 'w' or question == 'W' or question == 'wait' or question =='Wait':
        print "Starting server...",
        server(s)
        break
    elif question == 'c' or question == 'C' or question == 'connect' or question == 'Connect':
        client(s)
        break
    else:
        print "Sorry, I didn't quite get that.. little loud in here, you know."
        pass