1 files changed, 25 insertions, 17 deletions

diff --git a/src/main.rs b/src/main.rs
index c17f959..e0d0ed1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,3 +1,4 @@
+use std::env;
 use std::io::{self, Write};
 use std::sync::{Arc, Mutex};
 mod path;
@@ -9,30 +10,36 @@ use compose::Environment;
 /// Starts the main shell loop
 ///
 /// # Arguments
-/// * `prompt` - A string slice indicating the shell's prompt
-/// * `at_prompt` - A mutex, indicating whether or not user is at the prompt
+/// * `away` - A mutex, indicating whether or not user is at the prompt
+/// * `env` - The global shell state
 ///
 /// # Examples
 /// ```
 /// fn main() {
-///     let prompt = "|> ";
-///     let mut at_prompt = Arc::new(Mutex::new(false));
+///     let mut env = compose::env();
+///     let mut away = Arc::new(Mutex::new(false));
 ///     ...
-///     repl(prompt, &mut at_prompt);
+///     repl(&mut away, &mut env);
 /// }
 /// ```
-fn repl(prompt: &str, at_prompt: &mut Arc<Mutex<bool>>, env: &mut Environment) {
+fn repl(away: &mut Arc<Mutex<bool>>, env: &mut Environment) {
     // Initial path refresh on startup
     env.bins = path::refresh();
 
     // Main shell loop
     loop {
+        // Get the prompt
+        let prompt = match env::var("PS1") {
+            Ok(val) => val,
+            Err(_) => String::from("|> "),
+        };
+
         // Output the prompt
         print!("{}", prompt);
         io::stdout().flush().unwrap();
 
         // At the prompt
-        *at_prompt.lock().unwrap() = true;
+        *away.lock().unwrap() = false;
 
         // Wait for user input
         let mut poetry = String::new();
@@ -55,7 +62,7 @@ fn repl(prompt: &str, at_prompt: &mut Arc<Mutex<bool>>, env: &mut Environment) {
         }
 
         // Not at the prompt
-        *at_prompt.lock().unwrap() = false;
+        *away.lock().unwrap() = true;
 
         // Parse the poem
         let poem = Poem::read(poetry);
@@ -82,24 +89,25 @@ fn main() {
     // Compose the environment for dwvsh
     let mut env = compose::env();
 
-    // Set the prompt
-    let prompt = "|> ";
-    let mut at_prompt = Arc::new(Mutex::new(false));
-
     // Handle signals
+    let mut away = Arc::new(Mutex::new(true));
     unsafe {
-        let at_prompt = Arc::clone(&at_prompt);
+        let away = Arc::clone(&away);
         signal_hook::low_level::register(signal_hook::consts::SIGINT, move || {
-            if *at_prompt.lock().unwrap() {
+            if *away.lock().unwrap() {
+                println!();
+            } else {
+                let prompt = match env::var("PS1") {
+                    Ok(val) => val,
+                    Err(_) => String::from("|> "),
+                };
                 print!("\n{}", prompt);
                 io::stdout().flush().unwrap();
-            } else {
-                println!();
             }
         })
         .unwrap();
     };
 
     // Begin evaluating commands
-    repl(prompt, &mut at_prompt, &mut env);
+    repl(&mut away, &mut env);
 }