2 files changed, 62 insertions, 17 deletions

diff --git a/src/bin/server/database.rs b/src/bin/server/database.rs
index 9c51997..0fac387 100644
--- a/src/bin/server/database.rs
+++ b/src/bin/server/database.rs
@@ -1,22 +1,34 @@
 use rusqlite::Connection;
 
 pub fn init_db() -> Connection {
-    let conn = Connection::open("server.db").unwrap();
+    let db_path = std::env::var("SERVER_DB").unwrap_or_else(|_| "server.db".to_string());
+    let conn = Connection::open(&db_path).unwrap();
     conn.execute(
         "CREATE TABLE IF NOT EXISTS users (
             username TEXT PRIMARY KEY,
-            hashed_password TEXT NOT NULL
+            hashed_password TEXT NOT NULL,
+            account_id INTEGER UNIQUE,
+            balance INTEGER DEFAULT 1000
         )",
         [],
     ).unwrap();
+    // Add column if not exists (for migration)
+    let _ = conn.execute("ALTER TABLE users ADD COLUMN account_id INTEGER UNIQUE", []);
+    let _ = conn.execute("ALTER TABLE users ADD COLUMN balance INTEGER DEFAULT 1000", []);
     conn
 }
 
 pub fn add_user(conn: &Connection, username: &str, hashed_password: &str) -> Result<(), Box<dyn std::error::Error>> {
-    conn.execute(
-        "INSERT INTO users (username, hashed_password) VALUES (?1, ?2)",
-        [username, hashed_password],
-    )?;
+    loop {
+        let account_id: u32 = rand::random::<u32>();
+        match conn.execute(
+            "INSERT INTO users (username, hashed_password, account_id, balance) VALUES (?1, ?2, ?3, ?4)",
+            [username, hashed_password, &account_id.to_string(), "1000"],
+        ) {
+            Ok(_) => break,
+            Err(_) => continue, // Retry if account_id conflict
+        }
+    }
     Ok(())
 }
 
@@ -29,3 +41,23 @@ pub fn get_hashed_password(conn: &Connection, username: &str) -> Option<String>
         None
     }
 }
+
+pub fn get_account_id(conn: &Connection, username: &str) -> Option<u32> {
+    let mut stmt = conn.prepare("SELECT account_id FROM users WHERE username = ?1").unwrap();
+    let mut rows = stmt.query([username]).unwrap();
+    if let Ok(Some(row)) = rows.next() {
+        Some(row.get::<_, u32>(0).ok()?)
+    } else {
+        None
+    }
+}
+
+pub fn get_balance(conn: &Connection, username: &str) -> Option<i64> {
+    let mut stmt = conn.prepare("SELECT balance FROM users WHERE username = ?1").unwrap();
+    let mut rows = stmt.query([username]).unwrap();
+    if let Ok(Some(row)) = rows.next() {
+        Some(row.get::<_, i64>(0).ok()?)
+    } else {
+        None
+    }
+}
diff --git a/src/bin/server/handler.rs b/src/bin/server/handler.rs
index 2132d09..01fd9c9 100644
--- a/src/bin/server/handler.rs
+++ b/src/bin/server/handler.rs
@@ -19,6 +19,7 @@ pub fn handle_client(stream: TcpStream) {
         println!("[{} @ {}] Received: {}", user, client_addr, line);
         let response = process_message(&line, &mut user);
         writeln!(writer, "{}", response).unwrap();
+        println!("[{} @ {}] Sent: {}", user, client_addr, response);
     }
     println!("[{} @ {}] Client disconnected", user, client_addr);
 }
@@ -28,37 +29,49 @@ fn process_message(msg: &str, user: &mut String) -> String {
     match parts[0] {
         "REGISTER" => {
             if parts.len() != 3 {
-                "FAIL Invalid REGISTER format".to_string()
+                "ERROR Invalid REGISTER format".to_string()
             } else {
                 let username = parts[1];
-                let hashed_pass = parts[2];
+                let hashed_password = parts[2];
                 let db = DB_CONN.lock().unwrap();
-                match database::add_user(&db, username, hashed_pass) {
-                    Ok(_) => format!("OK Registered {}", username),
-                    Err(_) => "FAIL User already exists".to_string(),
+                match database::add_user(&db, username, hashed_password) {
+                    Ok(_) => format!("REGISTERED {}", username),
+                    Err(_) => "ERROR User already exists".to_string(),
                 }
             }
         }
         "AUTH" => {
             if parts.len() != 3 {
-                "FAIL Invalid AUTH format".to_string()
+                "ERROR Invalid AUTH format".to_string()
             } else {
                 let username = parts[1];
-                let provided_hash = parts[2];
+                let sent_hash = parts[2];
                 let db = DB_CONN.lock().unwrap();
                 if let Some(stored_hash) = database::get_hashed_password(&db, username) {
-                    if provided_hash == stored_hash {
+                    if sent_hash == stored_hash {
                         *user = username.to_string();
                         format!("OK {}", username)
                     } else {
-                        "FAIL Invalid credentials".to_string()
+                        "ERROR Invalid credentials".to_string()
                     }
                 } else {
-                    "FAIL User not found".to_string()
+                    "ERROR User not found".to_string()
+                }
+            }
+        }
+        "PING" => "PONG".to_string(),
+        "INFO" => {
+            if *user == "guest" {
+                "ERROR Not logged in".to_string()
+            } else {
+                let db = DB_CONN.lock().unwrap();
+                if let (Some(id), Some(balance)) = (database::get_account_id(&db, user), database::get_balance(&db, user)) {
+                    format!("INFO {} {} {}", user, id, balance)
+                } else {
+                    "ERROR Error retrieving user info".to_string()
                 }
             }
         }
-        "PING" => "PING".to_string(),
         _ => msg.to_string(), // Echo back for compatibility
     }
 }