2 files changed, 71 insertions, 17 deletions

diff --git a/src/openvic-simulation/scripts/ConditionalWeight.cpp b/src/openvic-simulation/scripts/ConditionalWeight.cpp
index ca42f1b..1bb83d0 100644
--- a/src/openvic-simulation/scripts/ConditionalWeight.cpp
+++ b/src/openvic-simulation/scripts/ConditionalWeight.cpp
@@ -26,23 +26,86 @@ 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;
 		}
 	);
 }
diff --git a/src/openvic-simulation/scripts/ConditionalWeight.hpp b/src/openvic-simulation/scripts/ConditionalWeight.hpp
index f667cd7..fcab0a6 100644
--- a/src/openvic-simulation/scripts/ConditionalWeight.hpp
+++ b/src/openvic-simulation/scripts/ConditionalWeight.hpp
@@ -14,7 +14,7 @@ namespace OpenVic {
 		using condition_weight_item_t = std::variant<condition_weight_t, condition_weight_group_t>;
 
 		enum class base_key_t : uint8_t {
-			BASE, FACTOR, MONTHS
+			BASE, FACTOR, TIME
 		};
 		using enum base_key_t;
 
@@ -31,15 +31,6 @@ namespace OpenVic {
 		ConditionalWeight(scope_type_t new_initial_scope, scope_type_t new_this_scope, scope_type_t new_from_scope);
 		ConditionalWeight(ConditionalWeight&&) = default;
 
-		static constexpr std::string_view base_key_to_string(base_key_t base_key) {
-			switch (base_key) {
-				case base_key_t::BASE: return "base";
-				case base_key_t::FACTOR: return "factor";
-				case base_key_t::MONTHS: return "months"; // TODO - add functionality for days or months or years
-				default: return "INVALID BASE KEY";
-			}
-		}
-
 		NodeTools::node_callback_t expect_conditional_weight(base_key_t base_key);
 
 		bool parse_scripts(DefinitionManager const& definition_manager);