feat: add organization prompt for template download

Cargo.toml

@@ -31,7 +31,7 @@ base32 = "0.5"
 base64 = "0.22"
 chrono = { version = "0.4", features = ["serde"] }
 clap = { version = "4.4", features = ["derive", "cargo", "color", "env"] }
-dialoguer = "0.11.0"
+dialoguer = { version = "0.11.0", features = ["fuzzy-select"] }
 dirs = "5.0"
 dunce = "1.0"
 fs4 = { version = "0.8", features = ["tokio"] }

src/commands/template.rs

@@ -1,9 +1,11 @@
 use crate::{
+    colors::ColorChoiceExt,
     commands::{
         install::{install, Install},
         login::check_login,
         GlobalArgs,
     },
+    dirs::pyproject_path,
     download::download_archive,
     error::{self, Result},
     git::init_repository,
@@ -25,6 +27,14 @@ use url::Url;
 )]
 pub struct GetCompetitionTemplate;
 
+#[derive(GraphQLQuery)]
+#[graphql(
+    query_path = "src/graphql/get_viewer_enabled_entities.graphql",
+    schema_path = "src/graphql/schema.graphql",
+    response_derives = "Debug,Clone"
+)]
+pub struct GetViewerEnabledEntities;
+
 #[derive(Args, Debug, Serialize)]
 #[command(author, version, about)]
 pub struct Template {
@@ -36,7 +46,7 @@ pub struct Template {
 
 pub async fn template(args: Template, global: GlobalArgs) -> Result<()> {
     let m = MultiProgress::new();
-    check_login(global.clone(), &m).await?;
+    let logged_in = check_login(global.clone(), &m).await?;
 
     let client = GraphQLClient::new(global.url.parse()?).await?;
 
@@ -63,7 +73,7 @@ pub async fn template(args: Template, global: GlobalArgs) -> Result<()> {
     pb.enable_steady_tick(std::time::Duration::from_millis(100));
     pb = m.add(pb);
 
-    let use_case = client
+    let competition = client
         .send::<GetCompetitionTemplate>(get_competition_template::Variables {
             slug: args.competition.clone(),
         })
@@ -74,16 +84,15 @@ pub async fn template(args: Template, global: GlobalArgs) -> Result<()> {
                 &format!("Competition '{}' not found", &args.competition),
                 "Please make sure the competition exists",
             )
-        })?
-        .use_case
-        .latest
-        .ok_or_else(|| {
-            error::system(
-                "No use case found",
-                "Please contact the competition organizer",
-            )
         })?;
 
+    let use_case = competition.use_case.latest.ok_or_else(|| {
+        error::system(
+            "No use case found",
+            "Please contact the competition organizer",
+        )
+    })?;
+
     let download_url = use_case
         .files
         .into_iter()
@@ -101,6 +110,54 @@ pub async fn template(args: Template, global: GlobalArgs) -> Result<()> {
         })?
         .download_url;
 
+    let organization = if logged_in {
+        let viewer = client
+            .send::<GetViewerEnabledEntities>(get_viewer_enabled_entities::Variables {
+                resource: competition.id,
+                action: get_viewer_enabled_entities::Action::CREATE_SUBMISSION_VERSION,
+            })
+            .await?
+            .viewer;
+        let viewer_orgs = viewer
+            .entities
+            .nodes
+            .iter()
+            .filter(|entity| entity.id != viewer.id)
+            .cloned()
+            .collect::<Vec<_>>();
+        if !viewer_orgs.is_empty() {
+            m.suspend(|| -> Result<_> {
+                let mut items = vec![format!("@{} (Myself)", viewer.username)];
+                items.extend(viewer_orgs.iter().map(|org| {
+                    format!("@{} ({})", org.username.clone(), org.display_name.clone())
+                }));
+                Result::Ok(
+                    dialoguer::FuzzySelect::with_theme(global.color.dialoguer().as_ref())
+                        .with_prompt("Would you like to submit with a team? (Press ESC to skip)")
+                        .items(&items)
+                        .interact_opt()
+                        .map_err(|err| {
+                            error::system(
+                                &format!("Could not select organization: {err}"),
+                                "Please try again",
+                            )
+                        })?
+                        .and_then(|index| {
+                            if index == 0 {
+                                None
+                            } else {
+                                viewer_orgs.into_iter().nth(index - 1)
+                            }
+                        }),
+                )
+            })?
+        } else {
+            None
+        }
+    } else {
+        None
+    };
+
     pb.set_message("Downloading competition template...");
     match download_archive(download_url, &destination, &pb).await {
         Ok(_) => {
@@ -123,6 +180,48 @@ pub async fn template(args: Template, global: GlobalArgs) -> Result<()> {
         }
     }
 
+    if let Some(organization) = organization {
+        let toml_path = pyproject_path(&destination);
+        let mut doc = tokio::fs::read_to_string(&toml_path)
+            .await
+            .map_err(|err| {
+                error::system(
+                    &format!("Failed to read {}: {err}", toml_path.display()),
+                    "Contact the competition organizer",
+                )
+            })?
+            .parse::<toml_edit::DocumentMut>()
+            .map_err(|err| {
+                error::system(
+                    &format!("Failed to parse {}: {err}", toml_path.display()),
+                    "Contact the competition organizer",
+                )
+            })?;
+        let aqora_config = doc
+            .get_mut("tool")
+            .and_then(|tool| tool.as_table_mut())
+            .and_then(|tool| tool.get_mut("aqora"))
+            .and_then(|aqora| aqora.as_table_mut())
+            .ok_or_else(|| {
+                error::system(
+                    &format!(
+                        "Failed to parse {}: Could not find tool.aqora",
+                        toml_path.display()
+                    ),
+                    "Contact the competition organizer",
+                )
+            })?;
+        aqora_config["entity"] = toml_edit::value(organization.username.clone());
+        tokio::fs::write(&toml_path, doc.to_string())
+            .await
+            .map_err(|err| {
+                error::system(
+                    &format!("Failed to write {}: {err}", toml_path.display()),
+                    "Check the permissions of the file",
+                )
+            })?;
+    }
+
     if !args.no_install {
         let install_global = GlobalArgs {
             project: destination.clone(),

src/graphql/get_viewer_enabled_entities.graphql

@@ -0,0 +1,15 @@
+query GetViewerEnabledEntities($resource: ID!, $action: Action!) {
+  viewer {
+    id
+    username
+    displayName
+    entities(permission: { to: $action, on: $resource }) {
+      nodes {
+        __typename
+        id
+        username
+        displayName
+      }
+    }
+  }
+}