import socket
import sys
import select
import os

BUFFER_SIZE = 1024
DEFAULT_HOST = 'localhost'
DEFAULT_PORT = 8181

def main():
    if len(sys.argv) != 3:
        SERVER_HOST = input(f"Enter server host (default: {DEFAULT_HOST}): ") or DEFAULT_HOST
        SERVER_PORT = input(f"Enter server port (default: {DEFAULT_PORT}): ") or DEFAULT_PORT
        SERVER_PORT = int(SERVER_PORT)
    else:
        SERVER_HOST = sys.argv[1]
        SERVER_PORT = int(sys.argv[2])

    # create a socket object
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        # connect to the server
        client_socket.connect((SERVER_HOST, SERVER_PORT))
        print(f"Connected to chat server at {SERVER_HOST}:{SERVER_PORT}")
        
        if os.name == 'nt':  # Windows-specific logic
            import msvcrt
            client_socket.setblocking(False)  # Set the socket to non-blocking mode

            while True:
                # Check for server messages
                try:
                    server_message = client_socket.recv(BUFFER_SIZE).decode('utf-8')
                    if server_message:
                        print(f"\n{server_message}")
                        print("> ", end="", flush=True)
                except BlockingIOError:
                    pass

                # Check for user input
                if msvcrt.kbhit():
                    client_message = input("> ")
                    if client_message.strip() == "%!":
                        print("Exiting chat...")
                        break
                    client_socket.send(client_message.encode('utf-8'))

        else:  # Linux/Unix-specific logic
            while True:
                # use select to wait for either server messages or user input
                read_sockets, _, _ = select.select([client_socket, sys.stdin], [], [])
                
                for sock in read_sockets:
                    if sock == client_socket:
                        # receive message from server
                        server_message = client_socket.recv(BUFFER_SIZE).decode('utf-8')
                        if not server_message:
                            print("Disconnected from server")
                            return
                        print(f"\n{server_message}")
                        print("> ", end="", flush=True)
                    else:
                        # prompt for user input
                        client_message = input()
                        if client_message.strip() == "%!":
                            print("Exiting chat...")
                            return
                        client_socket.send(client_message.encode('utf-8'))

    except Exception as e:
        print(f"Error: {e}")
    finally:
        client_socket.close()
        print("Connection closed")

if __name__ == "__main__":
    main()