diff --git a/src/main.rs b/src/main.rs
index 84942b9..af7891c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,18 +1,41 @@
 use std::io::{self, Write};
-use std::net::{TcpStream, ToSocketAddrs};
+use std::net::TcpStream;
 use std::time::Duration;
 use std::env;
-use colored::*;
+use colored::Colorize;
 use regex::Regex;
+use std::process;
 
 mod resp;
 
 const DEFAULT_PORT: u16 = 9880;
 const PROMPT_NAME: &str = "futriix";
-const CONNECTION_TIMEOUT: u64 = 5; // seconds
+const CONNECTION_TIMEOUT_SECS: u64 = 2;
 
 fn main() {
-    // Parse command line arguments
+    let (host, port) = parse_args();
+    let addr = format!("{}:{}", host, port);
+
+    let stream = match TcpStream::connect_timeout(
+        &addr.parse().unwrap(),
+        Duration::from_secs(CONNECTION_TIMEOUT_SECS)
+    ) {
+        Ok(stream) => stream,
+        Err(e) if e.kind() == io::ErrorKind::ConnectionRefused => {
+            eprintln!("{}", "Connection refused".red());
+            process::exit(1);
+        },
+        Err(e) => {
+            eprintln!("Connection error: {}", e);
+            process::exit(1);
+        }
+    };
+
+    println!("Connected to {}", addr.green());
+    run_repl_loop(stream, &host, port);
+}
+
+fn parse_args() -> (String, u16) {
     let args: Vec<String> = env::args().collect();
     let mut host = "127.0.0.1".to_string();
     let mut port = DEFAULT_PORT;
@@ -28,94 +51,63 @@ fn main() {
             }
         }
     }
+    (host, port)
+}
 
-    // Try to connect to server
-    let addr = format!("{}:{}", host, port);
-    let socket_addr = match format!("{}:{}", host, port).to_socket_addrs() {
-        Ok(mut addrs) => addrs.next().expect("No address found"),
-        Err(_e) => {
-            eprintln!("Failed to resolve address {}", addr);
-            return;
-        }
-    };
-
-    let stream = match TcpStream::connect_timeout(&socket_addr, Duration::from_secs(CONNECTION_TIMEOUT)) {
-        Ok(s) => s,
-        Err(_e) => {
-            eprintln!("Connection to port {} refused", port);
-            return;
-        }
-    };
-
-    println!("Connected to {}", addr);
-
-    // Main REPL loop
+fn run_repl_loop(stream: TcpStream, host: &str, port: u16) {
     let mut input = String::new();
     loop {
-        // Print prompt
-        print_prompt(&host, port);
-
-        // Read input
+        print_prompt(host, port);
         input.clear();
-        io::stdin().read_line(&mut input).expect("Failed to read input");
-        let input = input.trim();
+        
+        if io::stdin().read_line(&mut input).is_err() {
+            eprintln!("{}", "Failed to read input".red());
+            continue;
+        }
 
+        let input = input.trim();
         if input.is_empty() {
             continue;
         }
 
-        // Handle special commands
         if input.eq_ignore_ascii_case("quit") || input.eq_ignore_ascii_case("exit") {
             break;
         }
 
-        // Validate command before sending
         if !is_valid_command(input) {
-            eprintln!("Error: Invalid command format");
+            eprintln!("{}", "Error: Invalid command format".red());
             continue;
         }
 
-        // Send command to server
         match send_command(&stream, input) {
-            Ok(response) => {
-                print_response(&response);
-            }
+            Ok(response) => print_response(&response),
             Err(e) => {
-                eprintln!("Error: {}", e.to_string().replace("KeyDB", "Futriix"));
-                // Check if connection was lost
-                if e.kind() == io::ErrorKind::ConnectionAborted || 
-                   e.kind() == io::ErrorKind::ConnectionReset {
-                    eprintln!("Connection lost. Please restart the client.");
+                if is_connection_error(&e) {
+                    eprintln!("{}", "Connection error".red());
                     break;
                 }
+                eprintln!("{}", format!("Error: {}", e.to_string().replace("KeyDB", "Futriix")).red());
             }
         }
     }
 }
 
-fn is_valid_command(cmd: &str) -> bool {
-    if cmd.is_empty() {
-        return false;
-    }
-    
-    if cmd.chars().any(|c| c.is_control()) {
-        return false;
-    }
-    
-    cmd.split_whitespace().next().is_some()
-}
-
 fn print_prompt(host: &str, port: u16) {
     let prompt = format!("{}:{}:{}:~>", PROMPT_NAME, host, port);
     print!("{} ", prompt.green());
     io::stdout().flush().unwrap();
 }
 
+fn is_valid_command(cmd: &str) -> bool {
+    !cmd.is_empty() &&
+    !cmd.chars().any(|c| c.is_control()) &&
+    cmd.split_whitespace().next().is_some()
+}
+
 fn send_command(stream: &TcpStream, command: &str) -> io::Result<resp::Value> {
     let parts: Vec<&str> = command.split_whitespace().collect();
-    let mut resp_command = String::new();
+    let mut resp_command = format!("*{}\r\n", parts.len());
     
-    resp_command.push_str(&format!("*{}\r\n", parts.len()));
     for part in parts {
         resp_command.push_str(&format!("${}\r\n{}\r\n", part.len(), part));
     }
@@ -130,10 +122,10 @@ fn send_command(stream: &TcpStream, command: &str) -> io::Result<resp::Value> {
 fn print_response(value: &resp::Value) {
     match value {
         resp::Value::SimpleString(s) | resp::Value::BulkString(s) => {
-            println!("{}", s.replace("KeyDB", "Futriix"))
+            println!("{}", s.replace("KeyDB", "Futriix"));
         },
         resp::Value::Error(e) => {
-            println!("(error) {}", e.replace("KeyDB", "Futriix"))
+            println!("{}", format!("(error) {}", e.replace("KeyDB", "Futriix")).red());
         },
         resp::Value::Integer(i) => println!("(integer) {}", i),
         resp::Value::Array(arr) => {
@@ -141,17 +133,16 @@ fn print_response(value: &resp::Value) {
                 print!("{}) ", i + 1);
                 print_response(item);
             }
-        }
+        },
         resp::Value::Null => println!("(nil)"),
     }
 }
 
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_print_prompt() {
-        print_prompt("127.0.0.1", 9880);
-    }
+fn is_connection_error(error: &io::Error) -> bool {
+    matches!(
+        error.kind(),
+        io::ErrorKind::ConnectionAborted |
+        io::ErrorKind::ConnectionReset |
+        io::ErrorKind::BrokenPipe
+    )
 }
\ No newline at end of file