1 files changed, 31 insertions, 6 deletions

diff --git a/src/main.rs b/src/main.rs
index 4ebe0fa..e2e1b35 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,25 +1,36 @@
-use anyhow::{Context, Result};
+use anyhow::{bail, Context, Result};
 use rand::distributions::{Distribution, Uniform};
 use rand_word_gen::Model;
 use std::io::{BufWriter, Write};
 
+struct Config {
+    words: usize,
+    prefix: String,
+}
+
 fn main() -> Result<()> {
-    let words = parse_args()?;
+    let config = parse_args()?;
 
     let model = Model::new();
     let out = std::io::stdout();
     let mut out = BufWriter::new(out.lock());
     let mut rng = rand::thread_rng();
     let dist = Uniform::from(3..=6);
-    for _ in 0..words {
+    let first_char = config.prefix.bytes().last().map(|c| c as char);
+    for _ in 0..config.words {
         let len = dist.sample(&mut rng);
-        writeln!(out, "{}", model.generate(&mut rng, len))?;
+        writeln!(
+            out,
+            "{}{}",
+            config.prefix,
+            model.generate(&mut rng, len, first_char)
+        )?;
     }
 
     Ok(())
 }
 
-fn parse_args() -> Result<usize> {
+fn parse_args() -> Result<Config> {
     use clap::{crate_description, crate_version, value_t, App, Arg};
 
     let matches = App::new("rand-word-gen")
@@ -34,8 +45,22 @@ fn parse_args() -> Result<usize> {
                 .default_value("20")
                 .help("Sets number of generated words"),
         )
+        .arg(
+            Arg::with_name("prefix")
+                .short("p")
+                .long("prefix")
+                .value_name("PREFIX")
+                .default_value("")
+                .help("Sets prefix of generated words"),
+        )
         .get_matches();
 
     let words = value_t!(matches, "words", usize).context("'--words' must be a number")?;
-    Ok(words)
+    let mut prefix =
+        value_t!(matches, "prefix", String).context("Fatal error: failed to parse '--prefix'")?;
+    if !prefix.bytes().all(|c| c.is_ascii_alphabetic()) {
+        bail!("'--prefix' must be alphabetic");
+    }
+    prefix.make_ascii_lowercase();
+    Ok(Config { words, prefix })
 }