import socket
import threading
import rsa
from datetime import datetime


class Server:
    def __init__(self, host='0.0.0.0', port=5555):
        self.host = host
        self.port = port
        self.clients = {}
        self.public_keys = {}
        self.server_public_key = None
        self.server_private_key = None
        self.generate_keys()
        self.setup_server()

    def generate_keys(self):
        (self.server_public_key, self.server_private_key) = rsa.newkeys(2048)

    def setup_server(self):
        self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server.bind((self.host, self.port))
        self.server.listen()

    def broadcast(self, message, sender=None):
        for client in self.clients:
            if client != sender:
                try:
                    client_public_key = self.public_keys[client]
                    encrypted_msg = rsa.encrypt(message.encode(), client_public_key)
                    client.send(encrypted_msg)
                except:
                    self.remove_client(client)

    def handle_client(self, client):
        try:
            # Отправляем клиенту публичный ключ сервера
            client.send(self.server_public_key.save_pkcs1())

            # Получаем публичный ключ клиента
            client_public_key = rsa.PublicKey.load_pkcs1(client.recv(2048))
            self.public_keys[client] = client_public_key

            # Получаем имя клиента
            encrypted_name = client.recv(1024)
            client_name = rsa.decrypt(encrypted_name, self.server_private_key).decode()
            self.clients[client] = client_name

            print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {client_name} подключился.")
            self.broadcast(f"[СЕРВЕР] {client_name} присоединился к чату!", client)

            while True:
                try:
                    encrypted_msg = client.recv(4096)
                    if not encrypted_msg:
                        break

                    msg = rsa.decrypt(encrypted_msg, self.server_private_key).decode()

                    if msg.startswith('/file'):
                        self.handle_file_transfer(client, msg)
                    else:
                        self.broadcast(f"{self.clients[client]}: {msg}", client)

                except Exception as e:
                    print(f"[ОШИБКА] {e}")
                    break

        except Exception as e:
            print(f"[ОШИБКА] {e}")
        finally:
            self.remove_client(client)

    def handle_file_transfer(self, sender, command):
        try:
            parts = command.split(' ', 3)
            if len(parts) < 4:
                return

            recipient_name = parts[1]
            file_name = parts[2]
            file_data = parts[3].encode()

            # Находим получателя
            recipient = None
            for client in self.clients:
                if self.clients[client] == recipient_name:
                    recipient = client
                    break

            if recipient:
                # Отправляем уведомление получателю
                notification = f"/file_notification {self.clients[sender]} {file_name}"
                encrypted_notification = rsa.encrypt(notification.encode(), self.public_keys[recipient])
                recipient.send(encrypted_notification)

                # Ждем подтверждения
                response = rsa.decrypt(recipient.recv(1024), self.server_private_key).decode()
                if response == "/accept_file":
                    # Отправляем файл
                    encrypted_file_data = rsa.encrypt(file_data, self.public_keys[recipient])
                    recipient.send(encrypted_file_data)
                    sender.send(rsa.encrypt(f"[СЕРВЕР] Файл {file_name} отправлен {recipient_name}".encode(),
                                            self.public_keys[sender]))
                else:
                    sender.send(rsa.encrypt(f"[СЕРВЕР] {recipient_name} отказался принять файл".encode(),
                                            self.public_keys[sender]))
            else:
                sender.send(
                    rsa.encrypt(f"[СЕРВЕР] Пользователь {recipient_name} не найден".encode(), self.public_keys[sender]))

        except Exception as e:
            print(f"[ОШИБКА при передаче файла] {e}")

    def remove_client(self, client):
        if client in self.clients:
            name = self.clients[client]
            del self.clients[client]
            if client in self.public_keys:
                del self.public_keys[client]
            print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {name} отключился.")
            self.broadcast(f"[СЕРВЕР] {name} покинул чат.", client)
            client.close()

    def run(self):
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Сервер запущен на {self.host}:{self.port}")
        try:
            while True:
                client, address = self.server.accept()
                print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Новое соединение с {address[0]}:{address[1]}")

                thread = threading.Thread(target=self.handle_client, args=(client,))
                thread.start()
        except KeyboardInterrupt:
            print("\nСервер останавливается...")
        finally:
            self.server.close()


if __name__ == "__main__":
    server = Server()
    server.run()