From bf9108408ea587dc3708748c36794db48a8c6bc3 Mon Sep 17 00:00:00 2001
From: Paul Makles <me@insrt.uk>
Date: Wed, 1 Mar 2023 15:09:17 +0000
Subject: [PATCH] feat(safety): fetch additional context for snapshot

---
 .../delta/src/routes/safety/fetch_snapshot.rs | 66 ++++++++++++++++++-
 crates/quark/src/models/safety/snapshot.rs    | 19 +++++-
 2 files changed, 81 insertions(+), 4 deletions(-)

diff --git a/crates/delta/src/routes/safety/fetch_snapshot.rs b/crates/delta/src/routes/safety/fetch_snapshot.rs
index 2b46d5e9..81b8bef2 100644
--- a/crates/delta/src/routes/safety/fetch_snapshot.rs
+++ b/crates/delta/src/routes/safety/fetch_snapshot.rs
@@ -1,4 +1,7 @@
-use revolt_quark::models::{Snapshot, User};
+use std::collections::HashSet;
+
+use revolt_quark::models::snapshot::{SnapshotContent, SnapshotWithContext};
+use revolt_quark::models::{Channel, User};
 use revolt_quark::{Db, Error, Result};
 use rocket::serde::json::Json;
 
@@ -7,11 +10,68 @@ use rocket::serde::json::Json;
 /// Fetch a snapshot for a given report
 #[openapi(tag = "User Safety")]
 #[get("/snapshot/<report_id>")]
-pub async fn fetch_snapshot(db: &Db, user: User, report_id: String) -> Result<Json<Snapshot>> {
+pub async fn fetch_snapshot(
+    db: &Db,
+    user: User,
+    report_id: String,
+) -> Result<Json<SnapshotWithContext>> {
     // Must be privileged for this route
     if !user.privileged {
         return Err(Error::NotPrivileged);
     }
 
-    db.fetch_snapshot(&report_id).await.map(Json)
+    // Fetch snapshot
+    let snapshot = db.fetch_snapshot(&report_id).await?;
+
+    // Resolve and fetch IDs of associated content
+    let mut user_ids: HashSet<&str> = HashSet::new();
+    let mut channel_ids: HashSet<&str> = HashSet::new();
+
+    match &snapshot.content {
+        SnapshotContent::Message {
+            prior_context,
+            leading_context,
+            message,
+        } => {
+            for msg in prior_context {
+                user_ids.insert(&msg.author);
+            }
+
+            for msg in leading_context {
+                user_ids.insert(&msg.author);
+            }
+
+            user_ids.insert(&message.author);
+            channel_ids.insert(&message.channel);
+        }
+        _ => {
+            todo!()
+        }
+    }
+
+    // Collect user and channel IDs
+    let user_ids: Vec<String> = user_ids.into_iter().map(|s| s.to_owned()).collect();
+    let channel_ids: Vec<String> = channel_ids.into_iter().map(|s| s.to_owned()).collect();
+
+    // Fetch users and channels
+    let users = db.fetch_users(&user_ids).await?;
+    let channels = db.fetch_channels(&channel_ids).await?;
+
+    // Pull out first server from channels if possible
+    let server = if let Some(server_id) = channels.iter().find_map(|channel| match channel {
+        Channel::TextChannel { server, .. } => Some(server),
+        _ => None,
+    }) {
+        Some(db.fetch_server(server_id).await?)
+    } else {
+        None
+    };
+
+    // Return snapshot with context
+    Ok(Json(SnapshotWithContext {
+        snapshot,
+        users,
+        channels,
+        server,
+    }))
 }
diff --git a/crates/quark/src/models/safety/snapshot.rs b/crates/quark/src/models/safety/snapshot.rs
index e1d28df1..90eb7687 100644
--- a/crates/quark/src/models/safety/snapshot.rs
+++ b/crates/quark/src/models/safety/snapshot.rs
@@ -1,6 +1,6 @@
 use serde::{Deserialize, Serialize};
 
-use crate::models::{Message, Server, User};
+use crate::models::{Channel, Message, Server, User};
 
 /// Enum to map into different models
 /// that can be saved in a snapshot
@@ -35,3 +35,20 @@ pub struct Snapshot {
     /// Snapshot of content
     pub content: SnapshotContent,
 }
+
+/// Snapshot of some content with required data to render
+#[derive(Serialize, JsonSchema, Debug)]
+pub struct SnapshotWithContext {
+    /// Snapshot itself
+    #[serde(flatten)]
+    pub snapshot: Snapshot,
+    /// Users involved in snapshot
+    #[serde(rename = "_users", skip_serializing_if = "Vec::is_empty")]
+    pub users: Vec<User>,
+    /// Channels involved in snapshot
+    #[serde(rename = "_channels", skip_serializing_if = "Vec::is_empty")]
+    pub channels: Vec<Channel>,
+    /// Server involved in snapshot
+    #[serde(rename = "_server", skip_serializing_if = "Option::is_none")]
+    pub server: Option<Server>,
+}