1 files changed, 60 insertions, 7 deletions

diff --git a/src/openvic-simulation/scripts/ConditionalWeight.cpp b/src/openvic-simulation/scripts/ConditionalWeight.cpp
index ca42f1b..83385b2 100644
--- a/src/openvic-simulation/scripts/ConditionalWeight.cpp
+++ b/src/openvic-simulation/scripts/ConditionalWeight.cpp
@@ -26,23 +26,76 @@ static NodeCallback auto expect_modifier(
 }
 
 node_callback_t ConditionalWeight::expect_conditional_weight(base_key_t base_key) {
-	return expect_dictionary_keys(
-		// TODO - add days and years as options with a shared expected count of ONE_EXACTLY
-		base_key_to_string(base_key), ZERO_OR_ONE, expect_fixed_point(assign_variable_callback(base)),
-		"days", ZERO_OR_ONE, success_callback,
-		"years", ZERO_OR_ONE, success_callback,
+	key_map_t key_map;
+	bool successfully_set_up_base_keys = true;
+
+	switch (base_key) {
+	case BASE: {
+		successfully_set_up_base_keys &= add_key_map_entry(
+			key_map,
+			"base", ONE_EXACTLY, expect_fixed_point(assign_variable_callback(base))
+		);
+	} break;
+	case FACTOR: {
+		successfully_set_up_base_keys &= add_key_map_entry(
+			key_map,
+			"factor", ONE_EXACTLY, expect_fixed_point(assign_variable_callback(base))
+		);
+	} break;
+	case TIME: {
+		const auto time_callback = [this](std::string_view key, Timespan (*to_timespan)(Timespan::day_t)) -> auto {
+			return [this, key, to_timespan](uint32_t value) -> bool {
+				if (base == fixed_point_t::_0()) {
+					base = fixed_point_t::parse((*to_timespan)(value).to_int());
+					return true;
+				} else {
+					Logger::error(
+						"ConditionalWeight cannot have multiple base values - trying to set base to ", value, " ", key,
+						" when it already has a value equivalent to ", base, " days!"
+					);
+					return false;
+				}
+			};
+		};
+
+		successfully_set_up_base_keys &= add_key_map_entries(
+			key_map,
+			"days", ZERO_OR_ONE, expect_uint<uint32_t>(time_callback("days", Timespan::from_days)),
+			"months", ZERO_OR_ONE, expect_uint<uint32_t>(time_callback("months", Timespan::from_months)),
+			"years", ZERO_OR_ONE, expect_uint<uint32_t>(time_callback("years", Timespan::from_years))
+		);
+	} break;
+	default: {
+		successfully_set_up_base_keys = false;
+	}
+	}
+
+	if (!successfully_set_up_base_keys) {
+		return [base_key](ast::NodeCPtr node) -> bool {
+			Logger::error(
+				"Failed to set up base keys for ConditionalWeight with base value: ", static_cast<uint32_t>(base_key)
+			);
+			return false;
+		};
+	}
+
+	return expect_dictionary_key_map(
+		std::move(key_map),
 		"modifier", ZERO_OR_MORE, expect_modifier(condition_weight_items, initial_scope, this_scope, from_scope),
 		"group", ZERO_OR_MORE, [this](ast::NodeCPtr node) -> bool {
 			condition_weight_group_t items;
+
 			const bool ret = expect_dictionary_keys(
 				"modifier", ONE_OR_MORE, expect_modifier(items, initial_scope, this_scope, from_scope)
 			)(node);
+
 			if (!items.empty()) {
 				condition_weight_items.emplace_back(std::move(items));
 				return ret;
+			} else {
+				Logger::error("ConditionalWeight group must have at least one modifier!");
+				return false;
 			}
-			Logger::error("ConditionalWeight group must have at least one modifier!");
-			return false;
 		}
 	);
 }