import os
import socket
import threading
from Crypto.Cipher import AES

# Define AES encryption parameters
key = b"TheNeuralNineKey"
nonce = b"TheNeuralNineNce"
cipher = AES.new(key, AES.MODE_EAX, nonce)

IP = socket.gethostbyname(socket.gethostname())
PORT = 4466
ADDR = (IP, PORT)

SIZE = 4096
FORMAT = "utf-8"
DATA_PATH = "device1_data"

def handle_connection(conn, addr):
    print(f"[New connection] {addr} connected")
    conn.send("ok@Welcome to the file transfer system".encode(FORMAT))

    while True:
        data = conn.recv(SIZE).decode(FORMAT)
        cmd, *cmd_data = data.split("@")

        if cmd == "HELP":
            send_data = "ok@"
            send_data += "LIST: List all files\n"
            send_data += "SEND <filename>: Send a file\n"
            send_data += "RECEIVE <filename>: Receive a file\n"
            send_data += "DISCONNECT: Disconnect from the system\n"
            send_data += "HELP: List all commands"
            conn.send(send_data.encode(FORMAT))

        elif cmd == "DISCONNECT":
            break

        elif cmd == "LIST":
            files = os.listdir(DATA_PATH)
            send_data = "ok@"
            if not files:
                send_data += "The directory is empty"
            else:
                send_data += "\n".join(files)
            conn.send(send_data.encode(FORMAT))

        elif cmd == "SEND":
            filename = cmd_data[0]
            filepath = os.path.join(DATA_PATH, filename)

            if not os.path.exists(filepath):
                conn.send("ERROR@File not found".encode(FORMAT))
                continue

            with open(filepath, "rb") as f:
                file_data = f.read()

            encrypted_data = cipher.encrypt(file_data)
            file_size = len(encrypted_data)

            conn.send(f"READY@{file_size}".encode(FORMAT))
            conn.recv(SIZE)  # Wait for receiver's acknowledgment

            conn.sendall(encrypted_data)

            conn.recv(SIZE)  # Wait for receiver's confirmation

        elif cmd == "RECEIVE":
            filename = cmd_data[0]
            file_size = int(cmd_data[1])

            filepath = os.path.join(DATA_PATH, filename)

            conn.send("READY".encode(FORMAT))

            encrypted_data = b""
            remaining = file_size
            while remaining:
                chunk = conn.recv(min(SIZE, remaining))
                encrypted_data += chunk
                remaining -= len(chunk)

            decrypted_data = cipher.decrypt(encrypted_data)

            with open(filepath, "wb") as f:
                f.write(decrypted_data)

            send_data = "ok@File received successfully"
            conn.send(send_data.encode(FORMAT))

    print(f"[Disconnected] {addr} disconnected")

def start_server():
    print("[Starting] Server is starting")
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind(ADDR)
    server.listen()
    print("[Listening] Server is listening on", ADDR)

    while True:
        conn, addr = server.accept()
        thread = threading.Thread(target=handle_connection, args=(conn, addr))
        thread.start()

def connect_to_peer():
    peer_ip = input("Enter peer IP address: ")
    peer_port = int(input("Enter peer port: "))
    peer_addr = (peer_ip, peer_port)

    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client.connect(peer_addr)

    while True:
        data = client.recv(SIZE).decode(FORMAT)
        cmd, msg = data.split("@", 1)

        if cmd == "ok":
            print(f"{msg}")
        elif cmd == "ERROR":
            print(f"Error: {msg}")
        elif cmd == "READY":
            file_size = int(msg)
            client.send("ACK".encode(FORMAT))

            encrypted_data = b""
            remaining = file_size
            while remaining:
                chunk = client.recv(min(SIZE, remaining))
                encrypted_data += chunk
                remaining -= len(chunk)

            decrypted_data = cipher.decrypt(encrypted_data)
            filename = input("Enter the filename to save: ")
            filepath = os.path.join(DATA_PATH, filename)
            with open(filepath, "wb") as f:
                f.write(decrypted_data)
            print(f"File {filename} received successfully.")
            client.send("DONE".encode(FORMAT))

        user_input = input("> ")
        user_data = user_input.split(" ")
        cmd = user_data[0]

        if cmd == "HELP":
            client.send(cmd.encode(FORMAT))

        elif cmd == "DISCONNECT":
            client.send(cmd.encode(FORMAT))
            break

        elif cmd == "LIST":
            client.send(cmd.encode(FORMAT))

        elif cmd == "SEND":
            filename = user_data[1]
            filepath = os.path.join(DATA_PATH, filename)

            if not os.path.exists(filepath):
                print("Error: File not found")
                continue

            with open(filepath, "rb") as f:
                file_data = f.read()

            encrypted_data = cipher.encrypt(file_data)
            file_size = len(encrypted_data)

            client.send(f"RECEIVE@{filename}@{file_size}".encode(FORMAT))
            client.recv(SIZE)  # Wait for peer's READY signal

            client.sendall(encrypted_data)

        elif cmd == "RECEIVE":
            filename = user_data[1]
            client.send(f"SEND@{filename}".encode(FORMAT))

    print("Disconnected from the peer")
    client.close()

def main():
    os.makedirs(DATA_PATH, exist_ok=True)
    mode = input("Enter 's' for server mode or 'c' for client mode: ")
    if mode.lower() == 's':
        start_server()
    elif mode.lower() == 'c':
        connect_to_peer()
    else:
        print("Invalid mode. Exiting.")

if __name__ == "__main__":
    main()