import os
import socket
import select
import threading

CHATPORT = 8181
BUFFER_SIZE = 1024
MAX_CONNECT_QUEUE = 10

def format_file_info(file_path):
    """Helper function to format file information."""
    file_size = os.path.getsize(file_path)
    if file_size < 1000:
        size_str = f"{file_size} B"
    elif file_size < 1000000:
        size_str = f"{file_size / 1000:.2f} KB"
    else:
        size_str = f"{file_size / 1000000:.2f} MB"
    file_type = "Directory" if os.path.isdir(file_path) else "File"
    return f"{os.path.basename(file_path)} - {file_type} - {size_str}"

def handle_client(sock, client_address, socket_list, usernames, chat_rooms, client_states, user_paths):
    """Handle communication with a single client."""
    try:
        while True:
            message = sock.recv(BUFFER_SIZE).decode('utf-8')
            if not message:
                break

            state = client_states.get(sock, "NORMAL")
            if state == "WAITING_FOR_USERNAME":
                # Handle username setup
                usernames[sock] = message.strip()
                chat_rooms[sock] = "default"
                sock.send(b"\nUsername accepted. You are now in the 'default' chatroom.")
                client_states[sock] = "NORMAL"
                user_paths[sock] = None  # Initialize path for the user
            elif state == "NORMAL":
                if message.startswith("SET_PATH"):
                    # Handle setting the user's directory path
                    _, path = message.split(" ", 1)
                    if os.path.isdir(path):
                        user_paths[sock] = path
                        sock.send(b"Directory path set successfully.")
                    else:
                        sock.send(b"Error: Invalid directory path.")
                elif message.startswith("@") and "-ls" in message:
                    # Handle directory listing request
                    parts = message.split()
                    if len(parts) >= 2 and "-ls" in parts[1]:
                        usernames_part = parts[0]
                        if usernames_part.count("@") == 2:
                            receiving_username, sending_username = usernames_part[1:].split("@")
                            receiving_sock = next((s for s, u in usernames.items() if u == receiving_username), None)
                            if receiving_sock:
                                user_path = user_paths.get(receiving_sock)
                                if user_path and os.path.isdir(user_path):
                                    files = os.listdir(user_path)
                                    if files:
                                        file_list = "\n".join(
                                            [format_file_info(os.path.join(user_path, f)) for f in files]
                                        )
                                        sock.send(f"Files in {receiving_username}'s directory:\n{file_list}".encode('utf-8'))
                                    else:
                                        sock.send(f"No files in {receiving_username}'s directory.".encode('utf-8'))
                                else:
                                    sock.send(f"Error: {receiving_username} has no valid directory set.".encode('utf-8'))
                            else:
                                sock.send(b"Error: Receiving user not found.")
                        else:
                            sock.send(b"Error: Invalid message format.")
                elif message.startswith("@") and "-f" in message:
                    # Handle file request or transfer
                    parts = message.split()
                    if len(parts) >= 3 and "-f" in parts[1]:
                        usernames_part = parts[0]
                        if usernames_part.count("@") == 2:
                            receiving_username, sending_username = usernames_part[1:].split("@")
                            receiving_sock = next((s for s, u in usernames.items() if u == receiving_username), None)
                            if receiving_sock:
                                if len(parts) == 3:
                                    # File request
                                    filename = parts[2]
                                    receiving_sock.send(f"FILE_REQUEST {sending_username} {filename}".encode('utf-8'))
                                elif len(parts) > 5 and parts[-1] == "ENDFILE":
                                    # File transfer
                                    filename = parts[2]
                                    filetype = parts[3]
                                    filelength = int(parts[4])
                                    filebytes = " ".join(parts[5:-1]).encode('utf-8')
                                    receiving_sock.send(
                                        f"FILE_TRANSFER {sending_username} {filename} {filetype} {filelength}".encode('utf-8')
                                    )
                                    receiving_sock.send(filebytes)
                                    receiving_sock.send(b"ENDFILE")
                            else:
                                sock.send(b"Error: Receiving user not found.")
                        else:
                            sock.send(b"Error: Invalid message format.")
                elif message.strip() == "%!":
                    # Handle client exit
                    break
                else:
                    # Broadcast message to other clients in the same chat room
                    username = usernames.get(sock, "Unknown")
                    chat_room = chat_rooms.get(sock, "default")
                    send_message = f"{username} >> {message}".encode('utf-8')
                    for client_sock in socket_list:
                        if chat_rooms.get(client_sock) == chat_room:
                            try:
                                client_sock.send(send_message)
                            except Exception as e:
                                print(f"Error sending message to {client_sock}: {e}")
                                client_sock.close()
                                socket_list.remove(client_sock)
    except Exception as e:
        print(f"Error with client {client_address}: {e}")
    finally:
        # Clean up on client disconnect
        sock.close()
        socket_list.remove(sock)
        if sock in usernames:
            del usernames[sock]
        if sock in chat_rooms:
            del chat_rooms[sock]
        if sock in client_states:
            del client_states[sock]
        if sock in user_paths:
            del user_paths[sock]
        print(f"Client {client_address} disconnected.")

def main():
    # Initialize server
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.bind(('', CHATPORT))
    server_socket.listen(MAX_CONNECT_QUEUE)
    print(f"Chatserver listening on port: {CHATPORT}")

    # Shared data structures
    socket_list = []
    usernames = {}
    chat_rooms = {}
    client_states = {}
    user_paths = {}

    try:
        while True:
            client_socket, client_address = server_socket.accept()
            print(f"Accepted connection from {client_address}")
            socket_list.append(client_socket)

            # Initialize client state
            client_states[client_socket] = "WAITING_FOR_USERNAME"
            user_paths[client_socket] = None

            # Start a new thread for the client
            client_thread = threading.Thread(
                target=handle_client,
                args=(client_socket, client_address, socket_list, usernames, chat_rooms, client_states, user_paths),
                daemon=True
            )
            client_thread.start()
    except KeyboardInterrupt:
        print("\nShutting down server...")
    finally:
        server_socket.close()

if __name__ == "__main__":
    main()