import socket
import select

CHATPORT = 8181
BUFFER_SIZE = 1024
MAX_CONNECT_QUEUE = 10

def main():
    # set the flag variable to true for infinite loop
    quitFlag = True
    # create a socket object
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # set the socket options, SO_REUSEADDR allows to reuse the same port
    # SO_KEEPALIVE enables the socket to send keepalive messages
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
    # bind the socket to the address and port
    server_socket.bind(('', CHATPORT))
    # listen for incoming connections
    server_socket.listen(MAX_CONNECT_QUEUE)
    # set the server socket to non-blocking
    server_socket.setblocking(False)

    # create a list to store the sockets, server and client
    socket_list = [server_socket]
    # since a socket is associated with a username and a chatroom, 
    # and each one can only be associated one way at once, use dictionaries to
    # associate with each other.
    usernames = {}
    chat_rooms = {}
    # keep track of the states of each client connected, and if they are 
    # technically in an 'unsafe' state like in change username or chatroom mode.
    client_states = {}
    # let the user know the server is listening on port CHATPORT.
    print(f"Chatserver listening on port: {CHATPORT}")

    # infinite loop for the program
    while quitFlag:
        # use the select function to monitor the sockets
        read_sockets, _, _ = select.select(socket_list, [], [])
        # loop through the sockets
        for sock in read_sockets:
            # if the socket is the server socket, accept the connection
            if sock == server_socket:
                client_socket, client_address = server_socket.accept()
                client_socket.setblocking(False)
                socket_list.append(client_socket)
                print("\nAccepted connection from new client.")
                client_socket.send(b"\nWelcome to the chat server!\n")
                client_socket.send(b"\nUse '%!' to exit, or '%?' for help.\n")
                client_socket.send(b"\nType in a username: ")
                client_states[client_socket] = "WAITING_FOR_USERNAME"
            # otherwise the socket is a client socket.
            else:
                # try to receive the message from the client
                try:
                    message = sock.recv(BUFFER_SIZE).decode('utf-8')
                    if message:
                        state = client_states.get(sock, "NORMAL")
                        if state == "WAITING_FOR_USERNAME":
                            usernames[sock] = message.strip()
                            chat_rooms[sock] = "default"
                            sock.send(b"\nUsername accepted. You are now in the 'default' chatroom.")
                            client_states[sock] = "NORMAL"
                        elif state == "WAITING_FOR_NEW_USERNAME":
                            usernames[sock] = message.strip()
                            sock.send(b"\nUsername changed successfully.")
                            print(f"\nClient switched username to {usernames[sock]}.")
                            client_states[sock] = "NORMAL"
                        elif state == "WAITING_FOR_NEW_CHATROOM":
                            old_chatroom = chat_rooms[sock]
                            chat_rooms[sock] = message.strip()
                            sock.send(f"\nSwitched to chatroom '{chat_rooms[sock]}.'".encode('utf-8'))
                            print(f"\nClient switched to chatroom '{chat_rooms[sock]}.'")
                            if old_chatroom != "default" and old_chatroom not in chat_rooms.values():
                                print(f"\nChatroom '{old_chatroom}' is empty and removed.")
                            client_states[sock] = "NORMAL"
                        else:
                            if message.strip() == "%!":
                                sock.close()
                                socket_list.remove(sock)
                                del usernames[sock]
                                del chat_rooms[sock]
                                del client_states[sock]
                                print("\nClient has closed connection.")
                            elif message.strip() == "%?":
                                print("\nClient requested help.")
                                sock.send(b"Help: (%?)\n")
                                sock.send(b"Send messages like normal. Use enter to send a message.\n")
                                sock.send(b"Available Commands:\n")
                                sock.send(b"\t%? : Show this help prompt.\n")
                                sock.send(b"\t%! : Exit from the Chat Server.\n")
                                sock.send(b"\t%* : Switch to another chatroom.\n")
                                sock.send(b"\t%$ : Change Username.\n")
                            elif message.strip() == "%*":
                                sock.send(b"\nEnter chatroom name: ")
                                client_states[sock] = "WAITING_FOR_NEW_CHATROOM"
                            elif message.strip() == "%$":
                                sock.send(b"\nEnter new username: ")
                                client_states[sock] = "WAITING_FOR_NEW_USERNAME"
                            else:
                                username = usernames[sock]
                                chat_room = chat_rooms[sock]
                                send_message = f"{username} >> {message}".encode('utf-8')
                                for client_sock in socket_list:
                                    if client_sock != server_socket and chat_rooms.get(client_sock) == chat_room:
                                        client_sock.send(send_message)
                    else:
                        sock.close()
                        socket_list.remove(sock)
                        if sock in usernames:
                            del usernames[sock]
                        if sock in chat_rooms:
                            old_chatroom = chat_rooms[sock]
                            del chat_rooms[sock]
                            if old_chatroom != "default" and old_chatroom not in chat_rooms.values():
                                print(f"\nChatroom '{old_chatroom}' is empty and has been removed.")
                        if sock in client_states:
                            del client_states[sock]
                        print("\nClient has closed connection.")
                except Exception as err:
                    print(f"Error: {err}")
                    sock.close()
                    socket_list.remove(sock)
                    if sock in usernames:
                        del usernames[sock]
                    if sock in chat_rooms:
                        old_chatroom = chat_rooms[sock]
                        del chat_rooms[sock]
                        if old_chatroom != "default" and old_chatroom not in chat_rooms.values():
                            print(f"\nChatroom '{old_chatroom}' is empty and has been removed.")
                    if sock in client_states:
                        del client_states[sock]
                    print("\nClient has closed connection.")

    server_socket.close()

if __name__ == "__main__":
    main()