From 1b1f6a3eeea80736882013e82598857ac861811e Mon Sep 17 00:00:00 2001
From: Joey Hines <joey@ahines.net>
Date: Sun, 11 Feb 2024 16:37:46 -0700
Subject: [PATCH] Bot integration fixes

---
 Cargo.toml         |  2 +-
 src/api/mod.rs     | 35 ++++++++++++++++++++++++++++++-----
 src/api/models.rs  | 12 ++++++++----
 src/model/album.rs | 20 ++++++++++++--------
 4 files changed, 51 insertions(+), 18 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index fbdc7df..45dcce9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pic_ox"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/src/api/mod.rs b/src/api/mod.rs
index 34a2e65..69e2199 100644
--- a/src/api/mod.rs
+++ b/src/api/mod.rs
@@ -50,9 +50,18 @@ async fn get_album(
     Ok(Response(album))
 }
 
-async fn query_images(
-    image_query: Query<ImageQuery>,
+async fn get_image(
+    image_id: Path<u64>,
     State(context): State<PicContext>,
+) -> Result<Response<Image>, PicOxError> {
+    let image = context.db.get::<Image>(*image_id)?;
+
+    Ok(Response(image))
+}
+
+async fn query_images(
+    State(context): State<PicContext>,
+    Json(image_query): Json<ImageQuery>,
 ) -> Result<Response<Vec<Image>>, PicOxError> {
     let album_id = if let Some(album) = &image_query.album {
         Some(
@@ -62,6 +71,7 @@ async fn query_images(
                     album_name: Some(album.to_string()),
                 },
             )
+            .first()
             .ok_or(PicOxError::AlbumNotFound)?
             .id()
             .unwrap(),
@@ -150,8 +160,22 @@ async fn add_image(
     let user = headers.get("user").unwrap().to_str().unwrap();
     let user: u64 = user.parse().unwrap();
 
-    let mut album =
-        Album::find_album_by_query(&context.db, metadata.clone().unwrap().album).unwrap();
+    let album = Album::find_album_by_query(&context.db, metadata.clone().unwrap().album)
+        .first()
+        .cloned();
+
+    let mut album = if let Some(album) = album {
+        album
+    } else {
+        let name = metadata
+            .clone()
+            .unwrap()
+            .album
+            .album_name
+            .ok_or(PicOxError::AlbumNotFound)?;
+        let album = Album::new(&name, Vec::new(), user);
+        context.db.insert::<Album>(album)?
+    };
 
     let mut store_manager = context.store_manager.write().await;
     let img = store_manager
@@ -181,7 +205,7 @@ async fn add_image(
 async fn query_album(
     album_query: Query<AlbumQuery>,
     State(context): State<PicContext>,
-) -> Result<Response<Option<Album>>, PicOxError> {
+) -> Result<Response<Vec<Album>>, PicOxError> {
     let resp = Album::find_album_by_query(&context.db, album_query.0);
     Ok(Response(resp))
 }
@@ -233,6 +257,7 @@ pub async fn run_picox(db: Database, config: PicOxConfig) {
         .route("/api/album/:id", get(get_album))
         .route("/api/album/", get(query_album))
         .route("/api/image/", get(query_images))
+        .route("/api/image/:id", get(get_image))
         .with_state(context);
 
     let listener = tokio::net::TcpListener::bind(&config.host).await.unwrap();
diff --git a/src/api/models.rs b/src/api/models.rs
index 2f1838c..952b167 100644
--- a/src/api/models.rs
+++ b/src/api/models.rs
@@ -4,6 +4,7 @@ use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::Json;
 use axum_macros::FromRequest;
+use log::error;
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
 
@@ -92,10 +93,13 @@ impl IntoResponse for PicOxError {
                 StoreError::InvalidFile => (StatusCode::BAD_REQUEST, err.to_string()),
                 StoreError::OutOfStorage => (StatusCode::INSUFFICIENT_STORAGE, err.to_string()),
                 StoreError::ImageTooBig => (StatusCode::UNAUTHORIZED, err.to_string()),
-                StoreError::IOError(_) => (
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    "IO Error Has Occurred!".to_string(),
-                ),
+                StoreError::IOError(e) => {
+                    error!("IO error occurred handling request {}", e);
+                    (
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                        "IO Error Has Occurred!".to_string(),
+                    )
+                }
             },
             PicOxError::DbError(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
             PicOxError::AlbumNotFound => (
diff --git a/src/model/album.rs b/src/model/album.rs
index 2a3e869..2dc0795 100644
--- a/src/model/album.rs
+++ b/src/model/album.rs
@@ -29,14 +29,18 @@ impl Album {
         }
     }
 
-    pub fn find_album_by_query(db: &Database, album_query: AlbumQuery) -> Option<Self> {
-        if let Some(album_name) = &album_query.album_name {
-            db.filter(|_, album: &Album| album.album_name == *album_name)
-                .unwrap()
-                .next()
-        } else {
-            None
-        }
+    pub fn find_album_by_query(db: &Database, album_query: AlbumQuery) -> Vec<Self> {
+        db.filter(|_, album: &Album| {
+            if let Some(album_name) = &album_query.album_name {
+                if album_name != &album.album_name && !album.aliases.contains(album_name) {
+                    return false;
+                }
+            }
+
+            true
+        })
+        .unwrap()
+        .collect()
     }
 }