From 0f1c6edf0fa41cee490b511dd4fa0bdf9695f079 Mon Sep 17 00:00:00 2001
From: Profitroll <47523801+profitrollgame@users.noreply.github.com>
Date: Tue, 20 Dec 2022 23:59:47 +0100
Subject: [PATCH] Album covers added

---
 extensions/albums.py | 25 ++++++++++++++++++-------
 extensions/photos.py |  5 +++++
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/extensions/albums.py b/extensions/albums.py
index 5d3f0ce..b991286 100644
--- a/extensions/albums.py
+++ b/extensions/albums.py
@@ -30,7 +30,7 @@ async def album_create(name: str, title: str, current_user: User = Security(get_
 
     makedirs(path.join("data", "users", current_user.user, "albums", name), exist_ok=True)
 
-    uploaded = col_albums.insert_one( {"user": current_user.user, "name": name, "title": title} )
+    uploaded = col_albums.insert_one( {"user": current_user.user, "name": name, "title": title, "cover": None} )
 
     return UJSONResponse(
         {
@@ -52,7 +52,7 @@ async def album_find(q: str, current_user: User = Security(get_current_active_us
     return UJSONResponse(output)
 
 @app.patch("/albums/{id}", response_class=UJSONResponse, response_model=AlbumModified, description="Modify album's name or title by id")
-async def album_patch(id: str, name: Union[str, None] = None, title: Union[str, None] = None, current_user: User = Security(get_current_active_user, scopes=["albums.write"])):
+async def album_patch(id: str, name: Union[str, None] = None, title: Union[str, None] = None, cover: Union[str, None] = None, current_user: User = Security(get_current_active_user, scopes=["albums.write"])):
 
     try:
         album = col_albums.find_one( {"_id": ObjectId(id)} )
@@ -80,17 +80,24 @@ async def album_patch(id: str, name: Union[str, None] = None, title: Union[str,
     else:
         name = album["name"]
 
-    col_albums.update_one( {"_id": ObjectId(id)}, {"$set": {"name": name, "title": title}} )
+    if cover is not None:
+        image = col_photos.find_one( {"_id": ObjectId(cover), "album": album["name"]} )
+        cover = image["_id"].__str__() if image is not None else album["cover"]
+    else:
+        cover = album["cover"]
+
+    col_albums.update_one( {"_id": ObjectId(id)}, {"$set": {"name": name, "title": title, "cover": cover}} )
 
     return UJSONResponse(
         {
             "name": name,
-            "title": title
+            "title": title,
+            "cover": cover
         }
     )
 
 @app.put("/albums/{id}", response_class=UJSONResponse, response_model=AlbumModified, description="Modify album's name and title by id")
-async def album_put(id: str, name: str, title: str, current_user: User = Security(get_current_active_user, scopes=["albums.write"])):
+async def album_put(id: str, name: str, title: str, cover: str, current_user: User = Security(get_current_active_user, scopes=["albums.write"])):
 
     try:
         album = col_albums.find_one( {"_id": ObjectId(id)} )
@@ -108,18 +115,22 @@ async def album_put(id: str, name: str, title: str, current_user: User = Securit
     if 2 > len(title) > 40:
         return HTTPException(status_code=HTTP_406_NOT_ACCEPTABLE, detail="Album title must be >2 and <40 characters.")
 
+    image = col_photos.find_one( {"_id": ObjectId(cover), "album": album["name"]} )
+    cover = image["_id"].__str__() if image is not None else None # type: ignore
+    
     rename(
         path.join("data", "users", current_user.user, "albums", album["name"]),
         path.join("data", "users", current_user.user, "albums", name)
     )
 
     col_photos.update_many( {"user": current_user.user, "album": album["name"]}, {"$set": {"album": name}} )
-    col_albums.update_one( {"_id": ObjectId(id)}, {"$set": {"name": name, "title": title}} )
+    col_albums.update_one( {"_id": ObjectId(id)}, {"$set": {"name": name, "title": title, "cover": cover}} )
 
     return UJSONResponse(
         {
             "name": name,
-            "title": title
+            "title": title,
+            "cover": cover
         }
     )
 
diff --git a/extensions/photos.py b/extensions/photos.py
index 170e36e..0af1a84 100644
--- a/extensions/photos.py
+++ b/extensions/photos.py
@@ -114,6 +114,11 @@ async def photo_delete(id: str, current_user: User = Security(get_current_active
     except InvalidId:
         return HTTPException(status_code=HTTP_404_NOT_FOUND, detail="Could not find an image with such id.")
 
+    album = col_albums.find_one( {"name": image["album"]} )
+
+    if album is not None and album["cover"] == image["_id"].__str__():
+        col_albums.update_one( {"name": image["album"]}, {"$set": {"cover": None}} )
+
     remove(path.join("data", "users", current_user.user, "albums", image["album"], image["filename"]))
 
     return Response(status_code=HTTP_204_NO_CONTENT)