3 files changed, 93 insertions, 5 deletions

diff --git a/src/cli.rs b/src/cli.rs
index 38ac648..61b3d39 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -46,9 +46,21 @@ pub enum Commands {
         total: bool
     },
 
-    /// Add a new agent to lend or pay
-    Add {
+    /// Add, update or remove an agent
+    Agent {
         name: String,
+
+        /// Add a new agent (default)
+        #[arg(short,long)]
+        add: bool,
+
+        /// Delete a registered agent
+        #[arg(short,long, conflicts_with_all = ["add", "update"])]
+        delete: bool,
+
+        /// Update a registered agent
+        #[arg(short,long, conflicts_with_all = ["add", "delete"], value_name = "NEW_NAME")]
+        update: Option<String>
     },
 }
 
diff --git a/src/database.rs b/src/database.rs
index 8752255..da514ef 100644
--- a/src/database.rs
+++ b/src/database.rs
@@ -2,6 +2,8 @@ use rusqlite::{Connection, Result, Rows, Statement};
 use rusqlite::config::DbConfig::{*};
 use chrono::Utc;
 use std::path::PathBuf;
+use text_io::scan;
+use std::io::{stdout, Write};
 
 pub struct Database {
     date: i64,
@@ -23,6 +25,12 @@ impl Database {
     pub fn add_register(&self, filepath: &PathBuf) -> Result<()> {
         let conn = Connection::open(filepath)?;
         conn.set_db_config(SQLITE_DBCONFIG_ENABLE_FKEY, true)?;
+
+        if !&self.check_agent_exists(&conn)? {
+            println!("The Agent '{}' is not registered in the database", &self.person);
+            return Ok(());
+        }
+
         conn.execute("
             INSERT INTO Registers (agent_id, register_date, amount, note)
             VALUES (
@@ -45,6 +53,38 @@ impl Database {
         Ok(())
     }
 
+    pub fn update_agent(&self, filepath: &PathBuf, new_name: String) -> Result <()> {
+        let conn = Connection::open(filepath)?;
+
+        if !&self.check_agent_exists(&conn)? {
+            println!("The Agent '{}' is not registered in the database", &self.person);
+            return Ok(());
+        }
+
+        conn.execute("UPDATE Agents SET name=?1 WHERE name=?2", (&new_name, &self.person))?;
+
+        println!("Agent '{}' was updated to '{}'", &self.person, &new_name);
+        Ok(())
+    }
+
+    pub fn delete_agent(&self, filepath: &PathBuf) -> Result<()>{
+        let conn = Connection::open(filepath)?;
+        let prompt = format!("The Agent '{}' will be removed are you sure?", &self.person);
+
+        if !&self.check_agent_exists(&conn)? {
+            println!("The Agent '{}' is not registered in the database", &self.person);
+            return Ok(());
+        }
+        if !Database::ask_user_confirmaton(prompt.as_str()) {
+            println!("Operation aborted");
+            return Ok(());
+        }
+        conn.set_db_config(SQLITE_DBCONFIG_ENABLE_FKEY, true)?;
+        conn.execute("DELETE FROM Agents WHERE name=?1", (&self.person,))?;
+        println!("agent '{}' deleted successfully", &self.person);
+        Ok(())
+    }
+
     pub fn view_history(filepath: &PathBuf, filter: Option<String>) -> Result<()> {
         let conn = Connection::open(filepath)?;
         let mut hist_query: String = "
@@ -123,9 +163,37 @@ impl Database {
                 register_date INTEGER NOT NULL,
                 amount INTEGER NOT NULL,
                 note TEXT,
-                FOREIGN KEY(agent_id) REFERENCES Agents(id))", ())
+                FOREIGN KEY(agent_id) REFERENCES Agents(id)
+                    ON DELETE CASCADE)", ())
             .expect("SQL initialization error");
         println!("'{}' database created", &filepath.display());
         Ok(())
     }
+
+    fn check_agent_exists(&self, conn: &Connection) -> Result<bool> {
+        let mut stmt = conn.prepare("SELECT name FROM Agents WHERE name=?1")?;
+        let mut reg = stmt.query([&self.person])?;
+
+        if let Some(_) = reg.next()? {
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+
+    fn ask_user_confirmaton(prompt: &str) -> bool {
+        let mut input: String;
+        let complete_prompt = format!("{} (Y/n)", prompt);
+        loop {
+            print!("{}: ", complete_prompt);
+            stdout().flush().unwrap();
+            scan!("{}\n", input);
+
+            match input.to_lowercase().as_str() {
+                "y"|"" => return true,
+                "n" => return false,
+                _ => println!("Invalid input")
+            }
+        }
+    }
 }
diff --git a/src/main.rs b/src/main.rs
index be23ac3..8a4f338 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -26,6 +26,7 @@ fn main() -> Result<()> {
             let datapath = cli.database_path.unwrap();
             register.add_register(&datapath)?;
         },
+
         Commands::View {history, filter, total} => {
             let datapath = cli.database_path.expect("database path not found");
 
@@ -37,10 +38,17 @@ fn main() -> Result<()> {
             }
         },
 
-        Commands::Add {name} => {
+        Commands::Agent {name, add, delete, update} => {
             let register = Database::new(name, 0, None, false);
             let datapath = cli.database_path.unwrap();
-            register.add_agent(&datapath)?;
+
+            match (add, delete, update) {
+                (true, false, None) => register.add_agent(&datapath)?,
+                (false, true, None) => register.delete_agent(&datapath)?,
+                (false, false, None) => register.add_agent(&datapath)?,
+                (false, false, Some(new_name)) => register.update_agent(&datapath, new_name)?,
+                (_, _, _) => unreachable!()
+            }
         }
     };
     Ok(())