4 files changed, 42 insertions, 11 deletions

diff --git a/src/data/random_tables/henchmen.yaml b/src/data/random_tables/henchmen.yaml
index fd6c37f..f1a12f4 100644
--- a/src/data/random_tables/henchmen.yaml
+++ b/src/data/random_tables/henchmen.yaml
@@ -55,40 +55,40 @@ henchman_cleric_class:
   rows:
     - roll: 1
       steps:
-        - text: "Druid"
+        - class: druid
     - roll: 2-6
       steps:
-        - text: "Cleric"
+        - class: cleric
 
 henchman_fighter_class:
   formula: d10
   rows:
     - roll: 1
       steps:
-        - text: "Ranger"
+        - class: ranger
     - roll: 2
       steps:
-        - text: "Paladin"
+        - class: paladin
     - roll: 3-10
       steps:
-        - text: "Fighter"
+        - class: fighter
 
 henchman_magic_user_class:
   formula: d6
   rows:
     - roll: 1
       steps:
-        - text: "Illusionist"
+        - class: illusionist
     - roll: 2-6
       steps:
-        - text: "Magic-user"
+        - class: magic_user
 
 henchman_thief_class:
   formula: d6
   rows:
     - roll: 1
       steps:
-        - text: "Assassin"
+        - class: assassin
     - roll: 2-6
       steps:
-        - text: "Thief"
+        - class: thief
diff --git a/src/random_tables.rs b/src/random_tables.rs
index 8bc1e61..d059fd5 100644
--- a/src/random_tables.rs
+++ b/src/random_tables.rs
@@ -1,4 +1,5 @@
 use crate::dice;
+use crate::rules::classes::Classes;
 use include_dir::{include_dir, Dir};
 use serde::Deserialize;
 use serde_yaml;
@@ -20,8 +21,9 @@ struct TableRow {
 
 #[derive(Debug, Deserialize)]
 struct TableRowStep {
-    text: Option<String>,
+    class: Option<String>,
     table: Option<String>,
+    text: Option<String>,
 }
 
 pub struct RandomTables {
@@ -67,6 +69,13 @@ impl RandomTables {
                                 output_text.push_str(&inner_output);
                             }
 
+                            if let Some(class_string) = &step.class {
+                                let classes = Classes::new().unwrap();
+                                let class = classes.class(class_string)
+                                    .expect("Failed to load class.");
+                                output_text.push_str(&*class.name);
+                            }
+
                             output_text.push_str("\n");
                         }
                         return output_text.trim().replace("\n", ", ").to_string();
diff --git a/src/rules.rs b/src/rules.rs
index 9a925db..c1a85f3 100644
--- a/src/rules.rs
+++ b/src/rules.rs
@@ -1,2 +1,2 @@
 mod ability_scores;
-mod classes;
+pub(crate) mod classes;
diff --git a/src/rules/classes.rs b/src/rules/classes.rs
index 4d8f809..5cd29d6 100644
--- a/src/rules/classes.rs
+++ b/src/rules/classes.rs
@@ -1,10 +1,15 @@
 use crate::rules::ability_scores::AbilityScore;
 use serde::Deserialize;
+use serde_yaml;
+use std::collections::HashMap;
+use std::error::Error;
+use std::string::String;
 
 #[derive(Deserialize)]
 pub struct Class {
     pub name: String,
     pub prime_requisites: Vec<AbilityScore>,
+    #[serde(default)]
     pub npc_ability_score_modifiers: Vec<NpcAbilityScoreModifier>,
 }
 
@@ -13,3 +18,20 @@ pub struct NpcAbilityScoreModifier {
     pub ability_score: AbilityScore,
     pub modifier: i32,
 }
+
+const CLASSES_YAML: &str = include_str!("../data/rules/classes.yaml");
+
+pub struct Classes {
+    classes: HashMap<String, Class>,
+}
+
+impl Classes {
+    pub fn new() -> Result<Self, Box<dyn Error>> {
+        let classes: HashMap<String, Class> = serde_yaml::from_str(CLASSES_YAML)?;
+        Ok(Classes { classes })
+    }
+
+    pub fn class(&self, class_name: &str) -> Option<&Class> {
+        self.classes.get(class_name)
+    }
+}