From de9e668a72e0e00fe000bc0ee8a94b8a8c440a07 Mon Sep 17 00:00:00 2001
From: Kiana Sheibani <kiana.a.sheibani@gmail.com>
Date: Fri, 13 Oct 2023 19:13:17 -0400
Subject: [PATCH] Allow configuration of country and state when creating
 dataset

---
 src/main.rs | 34 ++++++++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index d7ad4d0..b6f1f2f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -148,11 +148,15 @@ fn dataset_new(name: Option<String>, auth_token: Option<String>) {
         .or_else(|| get_auth_token(&config_dir))
         .unwrap_or_else(|| error("Access token not provided", 1));
 
+    // Name
+
     let name = name.unwrap_or_else(|| {
         print!("Name of new dataset: ");
         read_string()
     });
 
+    // Game
+
     print!("Search games: ");
     let games = run_query::<VideogameSearch, _>(
         VideogameSearchVars {
@@ -184,6 +188,32 @@ fn dataset_new(name: Option<String>, auth_token: Option<String>) {
         name: game_name,
     } = games[index].clone();
 
+    // Location
+
+    print!("\nCountry to track ratings for (two-letter code, leave empty for none): ");
+    let country = {
+        let string = read_string();
+        if string.is_empty() {
+            None
+        } else {
+            Some(string)
+        }
+    };
+
+    let state = if country.as_ref().is_some_and(|s| s == "US" || s == "CA") {
+        print!("\nState/province to track ratings for (two-letter code, leave empty for none): ");
+        let string = read_string();
+        if string.is_empty() {
+            None
+        } else {
+            Some(string)
+        }
+    } else {
+        None
+    };
+
+    // Done configuring
+
     let connection =
         open_datasets(&config_dir).unwrap_or_else(|_| error("Could not open datasets file", 1));
     new_dataset(
@@ -193,8 +223,8 @@ fn dataset_new(name: Option<String>, auth_token: Option<String>) {
             last_sync: Timestamp(1),
             game_id,
             game_name,
-            country: None,
-            state: None,
+            country,
+            state,
             decay_rate: 0.5,
             period: (3600 * 24 * 30) as f64,
             tau: 0.2,