"""Chat Server for Client/Server programming lab in UC Berkeley's CS61A

Protocol Description:
    Logging In:
        Startup for 61AChat is a traditional three way handshake of the
        following format:
            Client -> Server "Client|Server|hello|"
            Server -> Client "Server|Client|welcome|"
            Client -> Server "Client|Server|thanks|"

    Messaging Another Client:
        To send a message to another client, one simply sends a "send-msg"
        message to the Server, which it forwards along to the correct
        recipient in a "receive-msg":
            Client1 -> Server "Client1|Client2|send-msg|message here"
            Server -> Client2 "Client1|Client2|receive-msg|message here"

    Logging Out:
        To log out, a client sends the logout message to the server:
            Client -> Server "Client|Server|logout|"
        At which point the Server removes the client from the table. In
        addition to removing the client upon request, we have the 

    Listing users:
        To get a list of users, the client sends the server a clients-list
        message. The server then replies with a receive-msg message with a body
        which lists all users.
            Client -> Server "Client|Server|clients-list|"
            Server -> Client "Server|Client|receive-msg|<list of clients>"

Written by Tom Magrino

Updated 4/8/2012
-Added kill-conn message to let server kick out clients
-Disallowed sending messages to yourself
-Send error message to client when recipient is unknown
"""

from ucb import main
from socket import socket, fromfd, AF_INET, SOCK_DGRAM
from socketserver import UDPServer, DatagramRequestHandler
from chatcommon import MSG_SIZE_LIMIT, decode_message, Message


class ChatServerHandler(DatagramRequestHandler):

    def handle(self):
        """Handle the new connection request to the server."""
        data, sock = self.request
        msg = decode_message(data.decode().strip())
        if msg.action == "hello":
            print("Hello from {0}. Handshaking!".format(msg.src))
            self.server.handshake(self.client_address, sock, msg)
        elif msg.action == "send-msg":
            print("Message {0}".format(str(msg)))
            if msg.src != msg.dst:
                self.server.send_msg(msg.src, msg.dst, msg.body)
        elif msg.action.startswith("_pig-"):        
            print("Message {0}".format(str(msg)))
            ### Fill this in to enable network pig! ###
        elif msg.action == "clients-list":
            print("Clients list request from {0}".format(str(msg)))
            self.server.send_client_list(msg.src)
        elif msg.action == "logout":
            print("{0} is logging off".format(msg.src))
            self.server.remove_client(msg.src)

class ChatServer(UDPServer):

    def __init__(self, server_address):
        """Start a new ChatServer."""
        UDPServer.__init__(self, server_address, ChatServerHandler)
        self.clients = {}

    def get_client(self, client):
        """Get the socket associated with this client."""
        return self.clients.get(client)

    def add_client(self, client, sock, addr):
        """Add a new client to the clients table, error if this client is
        already in the table!
        """
        if client in self.clients:
            body = "Error: another user has that name, quit and try again!"
            err_msg = Message("server", client, "receive-msg", body)
            sock.sendto(bytearray(str(err_msg).encode()), addr)
        else:
            self.clients[client] = sock, addr
            print(self.clients)

    def remove_client(self, client):
        """Remove a client from the clients table."""
        if client in self.clients:
            del self.clients[client]

    def send(self, msg):
        dest = self.get_client(msg.dst)
        src = self.get_client(msg.src)
        if dest:
            sock, addr = dest
            try:
                sock.sendto(bytearray(str(msg).encode()), addr)
            except:
                print("Had a problem sending {0}".format(msg))
        elif src:
            sock, addr = src
            
            body = "Unknown client: {0}".format(msg.dst)
            err_msg = Message("server", msg.src, "receive-msg", body)
            sock.sendto(bytearray(str(err_msg).encode()), addr)

            print("Client not in table: {0}".format(msg.dst))

    def send_msg(self, source, client, body):
        forward_msg = Message(source, client, "receive-msg", body)
        self.send(forward_msg)

    def send_client_list(self, client):
        body = "clients list:\n" + (", ".join(self.clients.keys()))
        client_list_msg = Message("server", client, "receive-msg", body)
        self.send(client_list_msg)

    def handshake(self, client_address, sock, msg):
        """Perform the three-way handshake to establish the connection."""
        client = msg.src
        welcome_msg = Message("server", client, "welcome", "")
        sock.sendto(bytearray(str(welcome_msg).encode()), client_address)
        buff = bytearray(b" " * MSG_SIZE_LIMIT)
        sock.recv_into(buff)
        reply = decode_message(buff.decode().strip())
        if reply.action == "thanks":
            self.add_client(client, sock, client_address)
        else:
            error_msg = "Error: Did not complete handshake!" 
            error_reply = Message("server", client, "receive-msg",
                                  error_msg)
            request.sendto(bytearray(str(error_reply).encode()), client_address)

    #For server errors/shutdown
    def kill_conn(self, client):
        print("Sending kill message to client {0}".format(client))
        kill_msg = Message("server", client, "kill-conn", "")
        self.send(kill_msg)
        self.remove_client(client)

@main
def run(*args):
    server = ChatServer(('', 0))
    print("Get your IP address by going to Google and searching 'ip'")
    print("Server port is: {0}".format(server.server_address[1]))
    print("Starting chat server, press <Control>-C to stop!")

    try:
        server.serve_forever()
    except:
        for client in list(server.clients.keys()):
            server.kill_conn(client)