Add new dependencies and implement chat UI with Ratatui

- Updated `Cargo.toml` to include new dependencies: `ratatui`, `crossterm`, `reqwest`, and `notify-rust`.
- Enhanced `main.rs` with a terminal-based chat UI using Ratatui.
- Added support for message rendering, input handling, and notifications.
- Implemented file download functionality and DM message handling.

Cargo.toml

@@ -16,3 +16,7 @@ rand = "0.8.5"
 rand_core = "0.6.4"
 crypto = "0.5.1"
 base64 = "0.21"
+ratatui = { version = "0.29.0", features = ["all-widgets"] }
+crossterm = "0.27"
+reqwest = "0.12.15"
+notify-rust = "4"

src/main.rs

@@ -5,11 +5,35 @@ use aes_gcm::{
 use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
 use colog;
 use log::{debug, error, info, warn};
+use notify_rust::Notification;
 use serde::Deserialize;
-use std::fs;
+use std::{
-use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
+    fs,
-use tokio::net::TcpStream;
+    fs::File,
+    io,
+    sync::{Arc, Mutex},
+};
+use tokio::{
+    io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader},
+    net::TcpStream,
+    sync::mpsc,
+    time::{sleep, Duration},
+};
 use x25519_dalek::{EphemeralSecret, PublicKey};
+// Ratatui imports
+use crossterm::{
+    event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode, KeyEventKind},
+    execute,
+    terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
+};
+use ratatui::{
+    backend::CrosstermBackend,
+    layout::{Constraint, Direction, Layout, Position},
+    style::{Color, Modifier, Style},
+    text::{Line, Span},
+    widgets::{Block, Borders, List, ListItem, Paragraph},
+    Terminal,
+};
 
 #[derive(Deserialize)]
 struct Config {
@@ -17,8 +41,45 @@ struct Config {
     port: Option<u16>,
 }
 
+// UI structs and enums
+enum InputMode {
+    Normal,
+    Editing,
+}
+
+struct ChatState {
+    input: String,
+    messages: Vec<(String, String)>, // (username, message)
+    input_mode: InputMode,
+    username: String,
+    should_quit: bool,
+}
+
+impl ChatState {
+    fn new(username: String) -> Self {
+        ChatState {
+            input: String::new(),
+            messages: Vec::new(),
+            input_mode: InputMode::Editing,
+            username,
+            should_quit: false,
+        }
+    }
+}
+
 #[tokio::main]
-async fn main() {
+async fn main() -> io::Result<()> {
+    let original_hook = std::panic::take_hook();
+    std::panic::set_hook(Box::new(move |panic_info| {
+        // Restore terminal
+        let _ = disable_raw_mode();
+        let mut stdout = io::stdout();
+        let _ = execute!(stdout, LeaveAlternateScreen, DisableMouseCapture);
+
+        // Call the original hook
+        original_hook(panic_info);
+    }));
+
     colog::init();
 
     let contents =
@@ -29,13 +90,11 @@ async fn main() {
     info!("Enter your username (or press Enter to use a random one): ");
     let mut input = String::new();
     std::io::stdin().read_line(&mut input).unwrap();
-    let mut username = input.trim().to_string();
+    let username = if input.trim().is_empty() {
-
+        format!("User{}", rand::random::<u32>())
-    if !username.is_empty() {
-        username = input.trim().to_string();
     } else {
-        username = format!("User{}", rand::random::<u32>());
+        input.trim().to_string()
-    }
+    };
 
     info!("Username: {}", username);
 
@@ -80,7 +139,7 @@ async fn main() {
         Ok(encrypted) => encrypted,
         Err(e) => {
             error!("Encryption error: {}", e);
-            return;
+            return Ok(());
         }
     };
 
@@ -88,34 +147,77 @@ async fn main() {
 
     if let Err(e) = writer.write_all((encoded + "\n").as_bytes()).await {
         error!("Failed to send username: {}", e);
-        return;
+        return Ok(());
     }
 
     info!("Starting the chat");
 
-    // Task for sending user input to the server
+    // Setup UI
-    let send_task = tokio::spawn(async move {
+    enable_raw_mode()?;
-        let stdin = tokio::io::stdin();
+    let mut stdout = io::stdout();
-        let mut stdin_reader = BufReader::new(stdin);
+    execute!(stdout, EnterAlternateScreen, EnableMouseCapture)?;
-        let mut input = String::new();
+    let backend = CrosstermBackend::new(stdout);
+    let mut terminal = Terminal::new(backend)?;
+
+    // Setup channels for communication
+    let (tx_ui, mut rx_ui) = mpsc::channel::<(String, String)>(100);
+    let (tx_net, mut rx_net) = mpsc::channel::<String>(100);
+
+    // Create shared state
+    let chat_state = Arc::new(Mutex::new(ChatState::new(username.clone())));
+    let chat_state_ui = Arc::clone(&chat_state);
+
+    // Task for UI handling
+    let ui_task = tokio::spawn(async move {
+        let mut chat_state = chat_state_ui;
 
         loop {
-            input.clear();
+            let should_quit = {
-            tokio::io::stdout().flush().await.unwrap();
+                let state = chat_state.lock().unwrap();
+                state.should_quit
+            };
 
-            stdin_reader.read_line(&mut input).await.unwrap();
+            if should_quit {
-
-            if input.trim().is_empty() {
-                continue;
-            }
-
-            if input.trim() == "/quit" {
-                info!("Disconnecting from server");
                 break;
             }
 
+            // Check for new messages from network
+            if let Ok(msg) = rx_ui.try_recv() {
+                let mut state = chat_state.lock().unwrap();
+                state.messages.push(msg);
+            }
+
+            // Handle input events
+            if let Ok(should_break) = ui_loop(&mut terminal, &mut chat_state, &tx_net) {
+                if should_break {
+                    break;
+                }
+            }
+
+            sleep(Duration::from_millis(16)).await; // ~60 fps refresh rate
+        }
+        if let Err(e) = disable_raw_mode() {
+            error!("Failed to disable raw mode: {}", e);
+        }
+
+        if let Err(e) = execute!(
+            terminal.backend_mut(),
+            LeaveAlternateScreen,
+            DisableMouseCapture
+        ) {
+            error!("Failed to leave alternate screen: {}", e);
+        }
+
+        if let Err(e) = terminal.show_cursor() {
+            error!("Failed to show cursor: {}", e);
+        }
+    });
+
+    // Task for sending messages to the server
+    let send_task = tokio::spawn(async move {
+        while let Some(input) = rx_net.recv().await {
             // Encrypt the input
-            let encrypted = match cipher_writer.encrypt(&nonce_writer, input.trim().as_bytes()) {
+            let encrypted = match cipher_writer.encrypt(&nonce_writer, input.as_bytes()) {
                 Ok(encrypted) => encrypted,
                 Err(e) => {
                     error!("Encryption error: {}", e);
@@ -141,6 +243,10 @@ async fn main() {
                 Ok(0) => {
                     // Server closed connection
                     info!("\nServer disconnected");
+                    tx_ui
+                        .send(("System".to_string(), "Server disconnected".to_string()))
+                        .await
+                        .ok();
                     break;
                 }
                 Ok(_) => {
@@ -168,12 +274,100 @@ async fn main() {
                         }
                     };
 
-                    info!("Received: {}", message.trim());
+                    if message.contains('|') {
-                    info!("Enter message: ");
+                        // Handle DM format
-                    tokio::io::stdout().flush().await.unwrap();
+                        let parts: Vec<&str> = message.splitn(2, '|').collect();
+                        if parts.len() == 2 {
+                            let sender = parts[0].trim();
+                            // The second part contains both receiver and message
+                            let receiver_and_message = parts[1].trim();
+                            // Split at the first space to separate receiver from message
+                            if let Some(space_pos) = receiver_and_message.find(' ') {
+                                let (receiver, content) = receiver_and_message.split_at(space_pos);
+                                if receiver != username {
+                                    // If the receiver is the same as the client, ignore
+                                    continue;
+                                }
+
+                                let content = content.trim_start();
+
+                                // Style as DM
+                                let dm_label = if sender == &username {
+                                    format!("DM to {}: ", receiver)
+                                } else {
+                                    format!("DM from {}: ", sender)
+                                };
+
+                                tx_ui
+                                    .send(("DM".to_string(), format!("{}{}", dm_label, content)))
+                                    .await
+                                    .ok();
+                            }
+                        }
+                    } else if message.contains("dl!") {
+                        // Handle file download request
+                        let parts: Vec<&str> = message.splitn(2, ' ').collect();
+                        if parts.len() == 2 {
+                            let filename = parts[1].trim();
+                            tx_ui
+                                .send((
+                                    "System".to_string(),
+                                    format!("Download request for file: {}", filename),
+                                ))
+                                .await
+                                .ok();
+                            let resp = reqwest::get(filename).await.expect("request failed");
+                            let body = resp.bytes().await.expect("body invalid");
+                            // get the file name from the end of the link
+                            let filename = filename.split('/').last().unwrap_or("file");
+                            // Create the file
+                            let mut out = File::create(filename).expect("failed to create file");
+                            let body_bytes = body.to_vec();
+                            io::copy(&mut &body_bytes[..], &mut out)
+                                .expect("failed to copy content");
+                            tx_ui
+                                .send((
+                                    "System".to_string(),
+                                    format!("Download completed, {}", filename),
+                                ))
+                                .await
+                                .ok();
+                        }
+                    } else if let Some(pos) = message.find(':') {
+                        let (sender, content) = message.split_at(pos);
+                        if sender == username {
+                            // If the sender is the same as the client, ignore
+                            continue;
+                        }
+
+                        // if the message contains a @username, highlight it
+                        if content.contains(&username) {
+                            // send the message in chat
+
+                            Notification::new()
+                                .summary("You got tagged in a message")
+                                .body(&format!("{}{}", sender, content))
+                                .show()
+                                .unwrap();
+                        }
+
+                        // Skip the colon and any space
+                        let content = content.trim_start_matches(|c| c == ':' || c == ' ');
+                        tx_ui
+                            .send((sender.to_string(), content.to_string()))
+                            .await
+                            .ok();
+                    } else {
+                        // If message format is different, treat as system message
+                        tx_ui.send(("System".to_string(), message)).await.ok();
+                    }
                 }
                 Err(e) => {
                     error!("Error reading from server: {}", e);
+                    tx_ui
+                        .send(("System".to_string(), format!("Error: {}", e)))
+                        .await
+                        .ok();
                     break;
                 }
             }
@@ -182,9 +376,135 @@ async fn main() {
 
     // Wait for tasks to complete
     tokio::select! {
+        _ = ui_task => (),
         _ = send_task => (),
         _ = receive_task => (),
     }
 
     info!("Client exiting");
+    Ok(())
+}
+
+// UI rendering function
+fn ui_loop<B: ratatui::backend::Backend>(
+    terminal: &mut Terminal<B>,
+    chat_state: &mut Arc<Mutex<ChatState>>,
+    tx_net: &mpsc::Sender<String>,
+) -> io::Result<bool> {
+    terminal.draw(|f| {
+        let size = f.area();
+
+        // Create layout with chat messages on top and input at bottom
+        let chunks = Layout::default()
+            .direction(Direction::Vertical)
+            .margin(1)
+            .constraints([Constraint::Min(3), Constraint::Length(3)])
+            .split(size);
+
+        let state = chat_state.lock().unwrap();
+
+        // Create messages list
+        let messages: Vec<ListItem> = state
+            .messages
+            .iter()
+            .map(|(username, message)| {
+                let username_style = if username == &state.username {
+                    Style::default().fg(Color::Green)
+                } else if username == "System" {
+                    Style::default().fg(Color::Yellow)
+                } else if username == "DM" {
+                    Style::default()
+                        .fg(Color::Magenta)
+                        .add_modifier(Modifier::BOLD)
+                } else {
+                    Style::default().fg(Color::Blue)
+                };
+                ListItem::new(Line::from(vec![
+                    Span::styled(format!("{}: ", username), username_style),
+                    Span::raw(message),
+                ]))
+            })
+            .collect();
+
+        let messages =
+            List::new(messages).block(Block::default().borders(Borders::ALL).title("Messages"));
+
+        // Input box
+        let input = Paragraph::new(state.input.as_str())
+            .style(match state.input_mode {
+                InputMode::Normal => Style::default(),
+                InputMode::Editing => Style::default().fg(Color::Yellow),
+            })
+            .block(Block::default().borders(Borders::ALL).title("Input"));
+
+        f.render_widget(messages, chunks[0]);
+        f.render_widget(input, chunks[1]);
+
+        // Set cursor position
+        match state.input_mode {
+            InputMode::Normal => {}
+            InputMode::Editing => {
+                f.set_cursor_position(Position::new(
+                    chunks[1].x + 1 + state.input.len() as u16,
+                    chunks[1].y + 1,
+                ));
+            }
+        }
+    })?;
+
+    // Handle events
+    if event::poll(Duration::from_millis(10))? {
+        if let Event::Key(key) = event::read()? {
+            if key.kind == KeyEventKind::Press {
+                let mut state = chat_state.lock().unwrap();
+
+                match state.input_mode {
+                    InputMode::Normal => match key.code {
+                        KeyCode::Char('e') => {
+                            state.input_mode = InputMode::Editing;
+                        }
+                        KeyCode::Char('q') => {
+                            state.should_quit = true;
+                            tx_net.try_send("/quit".to_string()).ok();
+                            return Ok(true);
+                        }
+                        _ => {}
+                    },
+                    InputMode::Editing => match key.code {
+                        KeyCode::Enter => {
+                            let message = state.input.drain(..).collect::<String>();
+                            if !message.is_empty() {
+                                drop(state); // Release mutex before async operation
+
+                                // Add message to UI
+                                let username_clone = {
+                                    let state = chat_state.lock().unwrap();
+                                    state.username.clone()
+                                };
+                                let mut state = chat_state.lock().unwrap();
+                                state
+                                    .messages
+                                    .push((username_clone.clone(), message.clone()));
+
+                                // Send to network
+                                tx_net.try_send(message).ok();
+                            }
+                        }
+                        KeyCode::Char(c) => {
+                            state.input.push(c);
+                        }
+                        KeyCode::Backspace => {
+                            state.input.pop();
+                        }
+                        KeyCode::Esc => {
+                            state.input_mode = InputMode::Normal;
+                        }
+                        _ => {}
+                    },
+                }
+            }
+        }
+    }
+
+    Ok(false)
 }