import os
import socket
import sys
import select

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])

    # Prompt user for a local directory path for file operations
    while True:
        local_path = input("Enter a local directory path for sending/receiving files: ").strip()
        if os.path.isdir(local_path):
            print(f"Using directory: {local_path}")
            break
        else:
            print("Invalid directory. Please enter a valid path.")

    # 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}")

        # Prompt user for a username and send it to the server
        username = input("Enter your username: ").strip()
        client_socket.send(username.encode('utf-8'))

        # Send the local path to the server
        client_socket.send(f"SET_PATH {local_path}".encode('utf-8'))

        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
                    if server_message.startswith("FILE_REQUEST"):
                        # Handle file request
                        _, sending_username, filename = server_message.split()
                        file_path = os.path.join(local_path, filename)
                        if os.path.isfile(file_path):
                            file_size = os.path.getsize(file_path)
                            with open(file_path, "rb") as f:
                                file_bytes = f.read()
                            filetype = "File"
                            client_socket.send(
                                f"@{sending_username}@CLIENT -f {filename} {filetype} {file_size} {file_bytes.decode('latin1')} ENDFILE".encode('utf-8')
                            )
                        else:
                            client_socket.send(
                                f"@{sending_username}@CLIENT -f {filename} NOTFOUND".encode('utf-8')
                            )
                    elif server_message.startswith("FILE_TRANSFER"):
                        # Handle file transfer
                        _, sending_username, filename, filetype, filelength = server_message.split()
                        file_bytes = client_socket.recv(int(filelength))
                        end_marker = client_socket.recv(BUFFER_SIZE).decode('utf-8')
                        if end_marker == "ENDFILE":
                            file_path = os.path.join(local_path, filename)
                            with open(file_path, "wb") as f:
                                f.write(file_bytes)
                            print(f"Received file '{filename}' from {sending_username}.")
                    else:
                        print(f"\n{server_message}")
                        print("> ", end="", flush=True)
                else:
                    # Prompt for user input
                    client_message = input("> ").strip()
                    if client_message == "%!":
                        print("Exiting chat...")
                        return

                    # Send message to server
                    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()