path: root/src/main.rs

use std::collections::HashMap;
use std::ffi::{OsStr, OsString};
use std::fmt;
use std::fs;
use std::io;
use std::path::PathBuf;
use std::process;

pub mod args;

use rand::{thread_rng, seq::{SliceRandom, IteratorRandom}};
use crate::args::{Args, ArgError, HELP_STR};

#[derive(Debug)]
enum Error {
    BadArguments(ArgError),
    NoSentences,
    MultipleSentences,
    BadTemplate(String),
    IOFail(io::Error),
    OutOfSentences,
    OutOfWords(OsString),
    UnknownCategory(OsString),
    NoCategories,
    DuplicateCategory(OsString),
    BadPath(PathBuf),
    EmptyOption,
    InvalidLiteral(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Error::BadArguments(e) =>
                write!(f, "Error while parsing arguments: {}\n{}", e, HELP_STR),
            Error::NoSentences =>
                write!(f, "no sentences.txt file found"),
            Error::MultipleSentences =>
                write!(f, "multiple sentences.txt files found, i'll be real i don't know how this is possible"),
            Error::BadTemplate(t) =>
                write!(f, "unable to parse template sentence: {}", t),
            Error::IOFail(e) =>
                write!(f, "IO error: {:?}", e),
            Error::OutOfSentences =>
                write!(f, "ran out of sentences"),
            Error::OutOfWords(c) =>
                write!(f, "ran out of words in category: {:?}", c),
            Error::UnknownCategory(c) =>
                write!(f, "tried to pick word in unknown category: {:?}", c),
            Error::NoCategories =>
                write!(f, "no categories provided"),
            Error::DuplicateCategory(c) =>
                write!(f, "tried to create a category twice: {:?}", c),
            Error::BadPath(p) =>
                write!(f, "failed to parse path: {:?}", p),
            Error::EmptyOption =>
                write!(f, "multi-option variable had an empty option"),
            Error::InvalidLiteral(s) =>
                write!(f, "literal option contains an inner quote-mark: {}", s),
        }
    }
}

impl From<io::Error> for Error {
    fn from(error: io::Error) -> Self {
        Error::IOFail(error)
    }
}

impl From<ArgError> for Error {
    fn from(error: ArgError) -> Self {
        Error::BadArguments(error)
    }
}

type ProgResult<T> = Result<T, Error>;

#[derive(Debug, Clone)]
enum DropIn {
    Any,
    Literal(String),
    Basic(String),
    Var(String),
    OneOf(Vec<DropIn>),
}

impl DropIn {
    pub fn parse(text: &str) -> ProgResult<Self> {
        if text.is_empty() {
            return Err(Error::EmptyOption)
        }

        if text == "?" {
            return Ok(DropIn::Any);
        }

        if text.contains('|') {
            let mut option_strs = text.split('|').collect::<Vec<&str>>();
            option_strs.dedup();
            let options = option_strs.into_iter()
                .map(DropIn::parse)
                .collect::<Result<Vec<_>, _>>()?;
            return Ok(DropIn::OneOf(options))
        }

        if let Some(lit) = text.strip_prefix('"') {
            if lit.contains('"') {
                return Err(Error::InvalidLiteral(text.to_string()));
            } else {
                return Ok(DropIn::Literal(lit.to_string()));
            }
        }

        if let Some(var_name) = text.strip_prefix('?') {
            return Ok(DropIn::Var(var_name.to_string()));
        }

        Ok(DropIn::Basic(text.to_string()))
    }
}

#[derive(Debug)]
struct CategorySet {
    var_table: HashMap<String, OsString>,
    categories: HashMap<OsString, Vec<String>>,
    unique: bool,
}

impl CategorySet {
    pub fn new(unique: bool) -> Self {
        CategorySet {
            var_table: HashMap::new(),
            categories: HashMap::new(),
            unique,
        }
    }

    fn resolve_variable(&mut self, var: &str) -> ProgResult<OsString> {
        if !self.var_table.contains_key(var) {
            let new_cat = self.random_category()?;
            self.var_table.insert(var.to_string(), new_cat);
        }
        Ok(self.var_table.get(var).unwrap().clone())
    }

    pub fn add_category(&mut self, name: OsString, mut words: Vec<String>) -> ProgResult<()> {
        words.retain(|s| !s.is_empty());
        words.shuffle(&mut thread_rng());
        match self.categories.insert(name.clone(), words) {
            None    => Ok(()),
            Some(_) => Err(Error::DuplicateCategory(name)),
        }
    }

    // not affected by --unique
    pub fn random_category(&self) -> ProgResult<OsString> {
        self.categories.keys()
            .choose(&mut thread_rng())
            .ok_or(Error::NoCategories)
            .map(|c| c.clone())
    }

    pub fn random_from_drop_in(&mut self, drop_in: &DropIn) -> ProgResult<String> {
        if let DropIn::Literal(s) = drop_in {
            return Ok(s.clone());
        }

        let simple_di = if let DropIn::OneOf(ds) = drop_in {
            ds.choose(&mut thread_rng())
                .expect("got a OneOf that is empty, somehow")
        } else {
            drop_in
        };

        let category = match simple_di {
            DropIn::Literal(s) => {
                return Ok(s.clone());
            },
            DropIn::Basic(s) => {
                OsString::from(s)
            },
            DropIn::Var(s) => {
                self.resolve_variable(s)?
            },
            DropIn::OneOf(_) => {
                panic!("Nested OneOf constructs are not supported")
            },
            DropIn::Any => {
                self.random_category()?
            },
        };

        let words = self.categories.get_mut(&category)
            .ok_or(Error::UnknownCategory(category.clone()))?;

        if self.unique {
            words.pop()
                .ok_or(Error::OutOfWords(category))
        } else {
            words.choose(&mut thread_rng())
                .ok_or(Error::OutOfWords(category))
                .map(|c| c.clone())
        }
    }
}

#[derive(Debug, Clone)]
struct Sentence {
    text_fragments: Vec<String>,
    drop_ins: Vec<DropIn>,
}

impl Sentence {
    pub fn parse(original: String) -> ProgResult<Self> {
        let mut text = original.as_str();
        let mut text_fragments = vec![];
        let mut drop_ins = vec![];

        while let Some((front, back)) = text.split_once("%%") {
            let (drop_in_body, rest) = back.split_once("%%")
                .ok_or(Error::BadTemplate(original.clone()))?;
            let drop_in = DropIn::parse(drop_in_body)?;

            text_fragments.push(front.to_string());
            drop_ins.push(drop_in);
            text = rest;
        }
        text_fragments.push(text.to_string());

        Ok(Sentence {
            text_fragments,
            drop_ins,
        })
    }

    pub fn generate(&self, categories: &mut CategorySet) -> ProgResult<String> {
        let mut text_iter = self.text_fragments.iter();
        let mut buffer = String::new();

        for drop_in in self.drop_ins.iter() {
            let text = text_iter.next()
                .expect("invariant: number of text = number of drop-in + 1");
            buffer.push_str(text);

            let drop_in_text = categories.random_from_drop_in(drop_in)?;
            buffer.push_str(drop_in_text.as_str());
        }

        let end_text = text_iter.next()
            .expect("invariant: number of text = number of drop-in + 1");
        buffer.push_str(end_text);

        Ok(buffer)
    }
}

#[derive(Debug)]
struct SentenceSet {
    sentences: Vec<Sentence>,
    unique: bool,
}

impl SentenceSet {
    pub fn new(mut sentence_strings: Vec<String>, unique: bool) -> ProgResult<Self> {
        sentence_strings.retain(|s| !s.is_empty());
        sentence_strings.shuffle(&mut thread_rng());

        let sentences = sentence_strings.into_iter()
            .map(Sentence::parse)
            .collect::<Result<Vec<_>, _>>()?;

        Ok(SentenceSet {
            sentences,
            unique,
        })
    }

    pub fn random_sentence(&mut self) -> ProgResult<Sentence> {
        if self.unique {
            self.sentences.pop()
                .ok_or(Error::OutOfSentences)
        } else {
            self.sentences.choose(&mut thread_rng())
                .ok_or(Error::OutOfSentences)
                .map(|s| s.clone())
        }
    }
}

fn read_lines(path: &PathBuf) -> ProgResult<Vec<String>> {
    let mut lines = vec![];

    let text = fs::read_to_string(path)?;
    for line in text.lines() {
        lines.push(line.to_string());
    }

    Ok(lines)
}

// REVIEW: There has to be a better way to do this. Separate into multiple functions?
fn read_files(dir_path: &str, unique: bool) -> ProgResult<(SentenceSet, CategorySet)> {
    let mut opt_sentences: Option<SentenceSet> = None;
    let mut categories = CategorySet::new(unique);

    let dir = fs::read_dir(dir_path)?;
    for entry in dir {
        let path_buf = entry?.path();
        if path_buf.is_file() {
            let extension = path_buf.extension()
                .ok_or(Error::BadPath(path_buf.clone()))?;
            if extension == OsStr::new("txt") {
                let lines = read_lines(&path_buf)?;
                let filename = path_buf.file_name()
                    .ok_or(Error::BadPath(path_buf.clone()))?;
                if filename == OsStr::new("sentences.txt") {
                    if opt_sentences.is_none() {
                        opt_sentences = Some(SentenceSet::new(lines, unique)?);
                    } else {
                        return Err(Error::MultipleSentences);
                    }
                } else {
                    let cat_name = path_buf.file_stem()
                        .expect("filename stem unwrapping osstr")
                        .to_os_string();
                    categories.add_category(cat_name, lines)?;
                }
            }
        }
    }

    match opt_sentences {
        None    => Err(Error::NoSentences),
        Some(s) => Ok((s, categories)),
    }
}

fn crash(e: Error) -> ! {
    eprintln!("[ERROR] {}", e);
    process::exit(1);
}

fn run() -> ProgResult<()> {
    let args = Args::parse()?;
    let (mut sentences, mut categories) = read_files(&args.directory, args.unique)?;

    for _i in 0..args.n_sentences {
        let sentence = sentences.random_sentence()?;
        let generated = sentence.generate(&mut categories)?;
        println!("{}", generated);
    }

    Ok(())
}

fn main() {
    if let Err(e) = run() {
        crash(e)
    }
}