client.py

"""
This file contains all the code for the graphical chat client.

The client is built using two major classes that work together:
- ChatUI : This is the main class. It's a Tkinter application that handles
  the GUI, manages the application's state (like chat history and who's
  online), and processes all user input.

- ClientNetworking : This is a focused helper class that ChatUI creates. Its
  only job is to manage the low-level socket connection to the server. It
  runs in a separate thread to keep the GUI from freezing and uses callback
  functions to pass server messages back to the ChatUI class for processing.
"""

import socket
import threading
import tkinter as tk
from tkinter import scrolledtext, messagebox
import sys
from datetime import datetime
import ctypes
import hamming
import rsa_module

# Try to set DPI awareness (Windows specific)
try:
    ctypes.windll.shcore.SetProcessDpiAwareness(1)
except AttributeError:
    pass  # Not on Windows or shcore not available


class ClientNetworking:
    """Handles raw socket communication, message framing, and threading."""

    """
    Handles the low-level networking stuff in a separate thread.

    Its whole purpose is to abstract away the raw socket communication so the
    main UI class doesn't have to deal with it. It connects to the server,
    sends properly framed messages, and listens for incoming data in a loop,
    maintaining a buffer and extracting complete | delimited messages.

    When it receives a message, it doesn't process it. Instead, it just
    calls the on payload received callback function, handing the data off
    to the ChatUI class to figure out what it means. It does the same thing
    for disconnects with the on disconnected callback.
    """

    def __init__(
        self, ip, port, name, public_key, on_payload_received, on_disconnected
    ):
        self.conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_addr = (ip, port)
        self.name = name
        self.public_key = public_key
        self.on_payload_received = on_payload_received
        self.on_disconnected = on_disconnected
        self.receive_buffer = ""
        self.is_connected = False

    def connect(self):
        try:
            self.conn.connect(self.server_addr)
            self.is_connected = True
            self.send_to_server(
                f"IAM:{self.name}:{self.public_key[0]}:{self.public_key[1]}"
            )
            threading.Thread(target=self.receive_loop, daemon=True).start()
            return True
        except Exception as e:
            self.on_disconnected(f"Connection failed: {e}")
            return False

    def receive_loop(self):

        # Maintain a buffer to handle messages that may arrive across multiple recv() calls
        # After each recv(), attempt to extract a complete '|' delimited message from the buffer
        # If successful, remove it from the buffer and process it
        try:
            while self.is_connected:
                data = self.conn.recv(4096)
                if not data:
                    break
                self.receive_buffer += data.decode()
                while "|" in self.receive_buffer:
                    try:
                        start = self.receive_buffer.index("|")
                        end = self.receive_buffer.index("|", start + 1)
                        payload = self.receive_buffer[start + 1 : end]
                        self.receive_buffer = self.receive_buffer[end + 1 :]

                        # Hand it off
                        self.on_payload_received(payload)

                    except ValueError:
                        break  # Incomplete message

        except (ConnectionResetError, ConnectionAbortedError):
            pass  # Disconnection is handled below
        except Exception as e:
            if self.is_connected:
                print(f"Receive error: {e}")  # Log unexpected errors
        finally:
            if self.is_connected:
                self.on_disconnected("Server closed connection.")
            self.close()

    def send_to_server(self, payload_content):
        if self.is_connected:
            try:
                self.conn.sendall(f"|{payload_content}|".encode())
            except Exception as e:
                self.on_disconnected(f"Send error: {e}")
                self.close()
        else:
            self.on_disconnected("Not connected. Cannot send message.")

    def close(self):
        if self.is_connected:
            self.is_connected = False
            try:
                self.conn.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass
            self.conn.close()


class ChatUI(tk.Tk):
    """
    Manages the entire graphical interface, application state, and user interaction.

    This class is the main driver of the client-side application. Its
    responsibilities can be broken down into a few key areas:

    1. Initialization and Setup:
       When a ChatUI object is created, it initializes the main Tkinter
       window, generates its own RSA keys for private messaging, and sets up
       all the GUI widgets in (_setup_ui). It then immediately creates and starts the
       ClientNetworking handler to connect to the server.

    2. UI and Display Logic:
       It's responsible for drawing everything the user sees. It uses a
       state variable, current_chat_target, to decide whether to show the
       main menu or a specific chat window. Functions like display_menu_or_chat 
       handle rendering the correct view and its history.

    3. State Management:
       It keeps track of all the important data in variables:
       - users_online: A dictionary of other users and their public keys.
       - chat_history: A dictionary mapping each user/group to a list of
         all its messages.
       - current_chat_target: The ID of the chat currently on screen.

    4. Handling User Input:
       When the user types and hits Enter, the handle input send function
       decides what to do. If in the menu, it treats the input as a menu
       choice. If in a chat, it prepares the message for sending—encrypting
       it with RSA for private chats or encoding it with Hamming for group
       chats—and then passes it to the network handler.

    5. Processing Server Messages:
       When the networking thread receives a message, it calls this class's
       process_server_payload function. This acts as a central dispatcher.
       It looks at the command (e.g., 'from', 'groupmsg', 'system:joined')
       and calls the appropriate handler function to update the app's state
       and refresh the GUI.
    """

    def __init__(self, ip, port, client_name):
        super().__init__()
        self.client_name = client_name
        self.title(f"Chat Client - {self.client_name}")
        self.geometry("400x375")
        self.resizable(False, False)
        self.config(bg="white")

        self.my_rsa_keys = rsa_module.generate_keys()  # Returns (e, d, n)
        print(
            f"My RSA Keys: e={self.my_rsa_keys[0]}, d={self.my_rsa_keys[1]}, n={self.my_rsa_keys[2]}"
        )

        # Keeps track of every other user currently connected to the server.
        # The key is the username, and the value is a dictionary holding their
        # public RSA key (e and n) and a counter for any unseen private messages.
        self.users_online = {}  # {username: {"e": pk_e, "n": pk_n, "unseen": 0}}
        
        # Stores the entire message history for all conversations.
        # The key is a chat identifier (either the special string 'group_chat'
        # or another user's name), and the value is a list of message dicts.
        # Each message dict contains : sender, text, timestamp, seen, type
        self.chat_history = {"group_chat": []}

        self.current_chat_target = None  # None for menu, "group_chat", or a username
        self.menu_map = {}  # Maps menu numbers to chat IDs

        self._setup_ui()

        # Define message handlers for commands from the server
        self.message_handlers = {
            "welcome": self._handle_welcome,
            "system": self._handle_system,
            "from": self._handle_private_message,
            "groupmsg": self._handle_group_message,
            "username_taken": self._handle_username_taken,
        }

        self.protocol("WM_DELETE_WINDOW", self.on_closing)
        self.bind("<Escape>", self.handle_escape_key)

        # Connect to server
        self.network_handler = ClientNetworking(
            ip,
            port,
            self.client_name,
            (self.my_rsa_keys[0], self.my_rsa_keys[2]),  # (e, n)
            lambda payload: self.after(0, self.process_server_payload, payload),
            lambda reason: self.after(0, self.process_disconnection, reason),
        )
        self.update_status(f"Connecting to {ip}:{port}...")
        if not self.network_handler.connect():
            self.after(100, self.display_menu_or_chat)  # Show disconnected state

    def _setup_ui(self):
        """Creates and places all Tkinter widgets."""

        self.display = scrolledtext.ScrolledText(
            self, state=tk.DISABLED, bg="white", wrap=tk.WORD, font=("Consolas", 9)
        )
        self.display.place(x=10, y=10, width=380, height=300)

        self.input_var = tk.StringVar()
        self.input_entry = tk.Entry(
            self,
            textvariable=self.input_var,
            font=("Consolas", 10),
            fg="#12e",
            bg="#ddf",
        )
        self.input_entry.place(x=10, y=320, width=300, height=25)
        self.input_entry.bind("<Return>", self.handle_input_send)

        self.send_button = tk.Button(
            self,
            text="Send",
            command=self.handle_input_send,
            font=("Consolas", 10),
            fg="#12e",
            bg="#ddf",
        )
        self.send_button.place(x=320, y=318, width=70, height=30)

        self.status_bar = tk.Label(
            self, text="Not connected.", bd=1, relief=tk.SUNKEN, anchor=tk.W
        )
        self.status_bar.place(x=0, y=355, relwidth=1.0, height=20)

    # --- UI Display Logic ---
    def display_menu_or_chat(self):
        """
        This is the main display router, called whenever the screen needs a full refresh.

        It checks the value of self.current chat target and, based on that,
        either calls display menu to show the user list or display chat to
        show the message history for a specific conversation.
        """

        self.display.config(state=tk.NORMAL)
        self.display.delete(1.0, tk.END)
        if self.current_chat_target is None:
            self._display_menu()
        else:
            self._display_chat()
        self.display.config(state=tk.DISABLED)
        self.input_entry.focus_set()

    def _display_menu(self):
        self._display_message_in_widget(
            "--- Main Menu ---\nEnter number to chat, or ESC to refresh.",
            is_system=True,
        )

        group_unseen = sum(
            1 for msg in self.chat_history.get("group_chat", []) if not msg["seen"]
        )
        self._display_message_in_widget(
            f"0. Group Chat" + (f" ({group_unseen} new)" if group_unseen else "")
        )
        self.menu_map = {0: "group_chat"}

        sorted_users = sorted(u for u in self.users_online if u != self.client_name)
        for i, username in enumerate(sorted_users, 1):
            user_unseen = self.users_online[username].get("unseen", 0)
            self._display_message_in_widget(
                f"{i}. {username}" + (f" ({user_unseen} new)" if user_unseen else "")
            )
            self.menu_map[i] = username

        if not sorted_users:
            if not self.network_handler or not self.network_handler.is_connected:
                self._display_message_in_widget(
                    "\nNot connected to server.", is_error=True
                )
            else:
                self._display_message_in_widget(
                    "\nNo other users online.", is_system=True
                )

    def _display_chat(self):
        chat_id = self.current_chat_target
        title = "Group Chat" if chat_id == "group_chat" else f"Chat with {chat_id}"
        self._display_message_in_widget(
            f"--- {title} ---\nPress ESC to return to main menu.", is_system=True
        )
        self._display_message_in_widget("-" * 30, is_system=True)

        for msg in self.chat_history.get(chat_id, []):
            display_text = f"[{msg['timestamp']}] "
            if msg["type"] == "system":
                self._display_message_in_widget(
                    display_text + msg["text"], is_system=True
                )
            else:
                self._display_message_in_widget(
                    f"{display_text}{msg['sender']}: {msg['text']}"
                )
            msg["seen"] = True

        if chat_id in self.users_online:
            self.users_online[chat_id]["unseen"] = 0

    # --- Input Handling ---
    def handle_input_send(self, event=None):
        """
        Triggered whenever the user presses Enter or clicks the Send button.

        This function acts as the primary handler for user input. It reads
        the text from the input box, and then, depending on whether the user
        is in the main menu or a chat, it passes the text to either
        handle menu selection or send chat message.
        """

        user_input = self.input_var.get().strip()
        self.input_var.set("")
        if not user_input or not self.network_handler.is_connected:
            if not self.network_handler.is_connected:
                self.update_status("Disconnected. Cannot send.")
            return

        user_input = user_input.replace("|", "").replace(":", "")
        if self.current_chat_target is None:
            self._handle_menu_selection(user_input)
        else:
            self._send_chat_message(user_input)

    def _handle_menu_selection(self, text):
        try:
            choice = int(text)
            if choice in self.menu_map:
                self.current_chat_target = self.menu_map[choice]
                self.display_menu_or_chat()
            else:
                self._display_message_in_widget(
                    "Invalid selection.", is_error=True, flash=True
                )
        except ValueError:
            self._display_message_in_widget(
                "Please enter a number.", is_error=True, flash=True
            )

    def _send_chat_message(self, text):
        """
        Prepares and sends a message to the currently selected chat target.

        This function is called by handle_input_send when the user is in a
        chat. It first adds the message to the local chat history so the user
        sees it immediately. Then, it checks if the target is the group chat
        (and applies Hamming encoding) or a private user (and applies RSA
        encryption) before sending the final payload to the server via the
        network handler.
        """

        target_id = self.current_chat_target
        self.add_to_history(target_id, "You", text, "sent")
        self._display_message_in_widget(
            f"[{datetime.now().strftime('%H:%M:%S')}] You: {text}"
        )

        if target_id == "group_chat":
            encoded_text, data_len = hamming.encode_string(text)
            total_b, data_b, _, _ = hamming.hamming_params(len(encoded_text))
            print(f"\n[SENDING] Encoded group message to '{encoded_text}'")
            print(
                f"          Using Hamming({total_b}, {data_b}) -> 2^{total_b - data_b} >= {data_b + (total_b - data_b) + 1}"
            )
            self.network_handler.send_to_server(f"groupsend:{data_len}:{encoded_text}")
        else:
            if target_id in self.users_online:
                recipient_pk = self.users_online[target_id]
                encrypted_msg = rsa_module.encrypt_string(
                    text, recipient_pk["e"], recipient_pk["n"]
                )
                self.network_handler.send_to_server(f"to:{target_id}:{encrypted_msg}")
            else:
                self._display_system_message(
                    f"User {target_id} is offline. Message not sent.", is_error=True
                )

    # --- Server Payload Processing ---


    def process_server_payload(self, payload):
        """
        The main dispatcher for all incoming server messages.

        The networking thread calls this function for every payload it
        receives. Its job is to extract the command from the message (like
        'from', 'groupmsg', etc.) and then call the appropriate handler
        function from the self.message handlers dictionary to deal with it.
        """

        command, _, data = payload.partition(":")
        handler = self.message_handlers.get(command, self._handle_unknown)
        handler(data or payload)

    def _handle_unknown(self, payload):
        self._display_message_in_widget(
            f"Unknown server message: {payload}", is_error=True
        )

    def _handle_username_taken(self, _):
        messagebox.showerror("Connection Error", "Username is already taken.")
        self.destroy()

    def _handle_welcome(self, data):
        welcome_text, _, users_part = data.partition("\nusers:")
        self.update_status(f"Connected! {welcome_text.splitlines()[0]}")
        self._display_system_message(welcome_text)
        if users_part:
            for entry in filter(None, users_part.split(",")):
                try:
                    name, e_str, n_str = entry.split("/")
                    self.users_online[name] = {
                        "e": int(e_str),
                        "n": int(n_str),
                        "unseen": 0,
                    }
                    self.chat_history.setdefault(name, [])
                except (ValueError, IndexError):
                    print(f"Malformed user entry: {entry}")
        self.display_menu_or_chat()

    def _handle_system(self, data):
        sys_type, _, sys_data = data.partition(":")
        if sys_type == "joined":
            name, pk_e, pk_n = sys_data.split(":")
            if name != self.client_name:
                self.users_online[name] = {"e": int(pk_e), "n": int(pk_n), "unseen": 0}
                self.chat_history.setdefault(name, [])
                self._display_system_message(f"{name} has joined.", flash=True)
                self.add_to_history(
                    "group_chat",
                    "System",
                    f"{name} has joined.",
                    "system",
                    seen=(self.current_chat_target == "group_chat"),
                )
                if self.current_chat_target is None:
                    self.display_menu_or_chat()
        elif sys_type == "left":
            name = sys_data
            if name in self.users_online:
                del self.users_online[name]
            if name in self.chat_history:
                del self.chat_history[name]
            self._display_system_message(f"{name} has left.", is_error=True, flash=True)
            self.add_to_history(
                "group_chat",
                "System",
                f"{name} has left.",
                "system",
                seen=(self.current_chat_target == "group_chat"),
            )
            if self.current_chat_target == name:
                self.current_chat_target = None
            self.display_menu_or_chat()
        elif sys_type == "shutdown":
            self._display_message_in_widget(f"SERVER ALERT: {sys_data}", is_error=True)
            self.process_disconnection("Server is shutting down.")
        else:  # notify, error
            self._display_system_message(sys_data, is_error=(sys_type == "error"))

    def _handle_private_message(self, data):
        """
        Handles an incoming private message from another user.

        It's called by the payload processor when a 'from' command is
        received. It uses this client's private RSA key to decrypt the
        message, adds the decrypted text to the correct chat history, and then
        updates the display. It also handles notification logic if the chat
        isn't currently open.
        """

        sender, _, encrypted_msg = data.partition(":")
        decrypted_msg = "[Decryption Failed]"
        try:
            decrypted_msg = rsa_module.decrypt_string(
                encrypted_msg, self.my_rsa_keys[1], self.my_rsa_keys[2]
            )
        except Exception as e:
            print(f"Decryption error from {sender}: {e}")

        is_current = self.current_chat_target == sender
        self.add_to_history(sender, sender, decrypted_msg, "received", seen=is_current)

        # If the same chat is being shown, show it on screen
        if is_current:
            self._display_message_in_widget(
                f"[{datetime.now().strftime('%H:%M:%S')}] {sender}: {decrypted_msg}",
                flash=True,
            )
        else:
            if sender in self.users_online:
                self.users_online[sender]["unseen"] += 1
            self._show_notification(f"New message from {sender}.")

    def _handle_group_message(self, data):
        """
        Handles an incoming group message from the server.

        It's called when a 'groupmsg' command is received. Its main task is
        to take the corrupted payload, pass it to the Hamming module's
        decode string function to fix the single-bit error, and then display
        the corrected message in the group chat window.
        """

        sender, data_len_str, encoded_msg = data.split(":", 2)
        msg = hamming.decode_string(encoded_msg, int(data_len_str))
        print(f"\n[RECEIVING] Group message from {sender}: '{encoded_msg}'")
        try:
            raw_decoded = hamming.bits_to_string(
                hamming.naive_extract(encoded_msg, int(data_len_str))
            )
            print(f"            Raw (no correction) : '{raw_decoded}'")
        except Exception:
            print("            Raw (no correction) : [Error displaying raw decode]")
        print(f"            After Hamming       : '{msg}'")

        is_current = self.current_chat_target == "group_chat"
        self.add_to_history("group_chat", sender, msg, "received", seen=is_current)
        if is_current:
            self._display_message_in_widget(
                f"[{datetime.now().strftime('%H:%M:%S')}] {sender} : {msg}", flash=True
            )
        else:
            self._show_notification(f"New group message from {sender}.")


    # --- Utility and Helper Methods ---

    def add_to_history(self, chat_id, sender, text, msg_type, seen=True):
        self.chat_history.setdefault(chat_id, []).append(
            {
                "sender": sender,
                "text": text,
                "timestamp": datetime.now().strftime("%H:%M:%S"),
                "seen": seen,
                "type": msg_type,
            }
        )

    def _display_message_in_widget(
        self, msg, is_system=False, is_error=False, flash=False
    ):
        self.display.config(state=tk.NORMAL)
        tags = ()
        if is_system:
            self.display.tag_config("system", foreground="blue")
            tags += ("system",)
        if is_error:
            self.display.tag_config(
                "error", foreground="red", font=("Consolas", 9, "bold")
            )
            tags += ("error",)
        self.display.insert(tk.END, msg + "\n", tags)
        self.display.yview(tk.END)
        self.display.config(state=tk.DISABLED)
        if flash:
            original_bg = self.display.cget("background")
            flash_color = "#fdd" if is_error else "#dfd"
            self.display.config(bg=flash_color)
            self.after(400, lambda: self.display.config(bg=original_bg))

    def _display_system_message(self, text, is_error=False, flash=False):
        prefix = (
            "SYSTEM: "
            if self.current_chat_target is None
            else "[Notification] System: "
        )
        self._display_message_in_widget(
            prefix + text, is_system=True, is_error=is_error, flash=flash
        )

    def _show_notification(self, text):
        if self.current_chat_target is None:
            self.display_menu_or_chat()
        self._display_message_in_widget(
            f"[Notification] {text}", is_system=True, flash=True
        )

    def process_disconnection(self, reason):
        """
        This function is the callback for when the server connection is lost.

        Its purpose is to update the GUI to reflect the disconnected state.
        It updates the status bar with the reason for disconnection and
        disables the message input and send buttons so the user can't try
        to send anything while offline.
        """
                
        self.update_status(f"Disconnected: {reason}")
        self.input_entry.config(state=tk.DISABLED)
        self.send_button.config(state=tk.DISABLED)
        self.display.config(bg="#eee")
        if self.current_chat_target is None:
            self.display_menu_or_chat()
        else:
            self._display_system_message(reason, is_error=True)

    def handle_escape_key(self, event=None):
        self.current_chat_target = None
        self.display_menu_or_chat()

    def update_status(self, text):
        self.status_bar.config(text=text)

    def on_closing(self):
        if messagebox.askokcancel("Quit", "Do you want to quit?"):
            if self.network_handler:
                self.network_handler.close()
            self.destroy()


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python client.py IP:Port YourName")
        sys.exit(1)

    try:
        ip_str, port_str = sys.argv[1].split(":")
        port_num = int(port_str)
        client_name = sys.argv[2]

        if any(c in client_name for c in "|:") or not (1 <= len(client_name) <= 20):
            raise ValueError(
                "Name must be 1-20 characters and cannot contain '|' or ':'."
            )

        app = ChatUI(ip_str, port_num, client_name)
        app.mainloop()

    except ValueError as e:
        print(f"Error: Invalid arguments. {e}")
        sys.exit(1)