Update app.py

app.py

@@ -147,6 +147,14 @@ MONGODB_URL = os.getenv("MONGODB_URL")
 client = None
 database = None
 
+# --------------------- Admin Panel DB (categories + media_clicks) ---------------------
+ADMIN_MONGO_URL = os.getenv("ADMIN_MONGO_URL")
+admin_client = AsyncIOMotorClient(ADMIN_MONGO_URL)
+admin_db = admin_client.adminPanel
+
+categories_col = admin_db.categories
+media_clicks_col = admin_db.media_clicks
+
 # --------------------- FastAPI ---------------------
 fastapi_app = FastAPI()
 
@@ -285,69 +293,181 @@ async def health():
 async def face_swap_api(
     source: UploadFile = File(...),
     target_category_id: str = Form(...),
+    user_id: Optional[str] = Form(None),
+    category_id: Optional[str] = Form(None),
     user_email: str = Depends(verify_firebase_token)
 ):
-    # start_time = datetime.utcnow()
     start_time = datetime.now(timezone.utc)
+
     try:
+        # READ SOURCE
         src_bytes = await source.read()
         src_key = f"bikini-theme/source/{uuid.uuid4().hex}_{source.filename}"
         upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
 
+        # ------------------------------------------------------------------
+        # OPTIONAL CATEGORY CLICK LOGGING
+        # ------------------------------------------------------------------
+        if user_id and category_id:
+            try:
+                user_oid = ObjectId(user_id)
+                category_oid = ObjectId(category_id)
+                now = datetime.utcnow()
+
+                # check category exists
+                category_doc = await categories_col.find_one({"_id": category_oid})
+                if category_doc:
+                    update_res = await media_clicks_col.update_one(
+                        {"userId": user_oid, "categories.categoryId": category_oid},
+                        {
+                            "$set": {
+                                "updatedAt": now,
+                                "categories.$.lastClickedAt": now
+                            },
+                            "$inc": {"categories.$.click_count": 1}
+                        }
+                    )
+
+                    if update_res.matched_count == 0:
+                        await media_clicks_col.update_one(
+                            {"userId": user_oid},
+                            {
+                                "$setOnInsert": {"createdAt": now},
+                                "$set": {"updatedAt": now},
+                                "$push": {
+                                    "categories": {
+                                        "categoryId": category_oid,
+                                        "click_count": 1,
+                                        "lastClickedAt": now
+                                    }
+                                }
+                            },
+                            upsert=True
+                        )
+                else:
+                    logger.error(f"Invalid category_id (not in categories collection): {category_id}")
+
+            except Exception as log_err:
+                logger.error(f"MEDIA_CLICK ERROR: {log_err}")
+
+        # ------------------------------------------------------------------
+        # TARGET IMAGE DOWNLOAD
+        # ------------------------------------------------------------------
         target_filename = f"{target_category_id}.png"
         target_url = f"https://{DO_SPACES_BUCKET}.{DO_SPACES_REGION}.digitaloceanspaces.com/bikini-theme/target/{target_filename}"
 
         resp = requests.get(target_url)
         if resp.status_code != 200:
-            # end_time = datetime.utcnow()
             end_time = datetime.now(timezone.utc)
-            await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
-            raise HTTPException(status_code=404, detail=f"Target image not found at {target_url}")
+            await log_faceswap_hit(user_email, "error", start_time, end_time)
+            raise HTTPException(404, f"Target not found at {target_url}")
 
         tgt_bytes = resp.content
 
-        src_array = np.frombuffer(src_bytes, np.uint8)
-        tgt_array = np.frombuffer(tgt_bytes, np.uint8)
-        src_bgr = cv2.imdecode(src_array, cv2.IMREAD_COLOR)
-        tgt_bgr = cv2.imdecode(tgt_array, cv2.IMREAD_COLOR)
+        # FACE SWAP
+        src_bgr = cv2.imdecode(np.frombuffer(src_bytes, np.uint8), cv2.IMREAD_COLOR)
+        tgt_bgr = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
 
         if src_bgr is None or tgt_bgr is None:
-            #end_time = datetime.utcnow()
             end_time = datetime.now(timezone.utc)
-
-            await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
-            raise HTTPException(status_code=400, detail="Invalid image data")
+            await log_faceswap_hit(user_email, "error", start_time, end_time)
+            raise HTTPException(400, "Invalid image data")
 
         src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
         tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
 
         final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
         if err:
-            end_time = datetime.utcnow()
-            await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
-            raise HTTPException(status_code=500, detail=err)
+            end_time = datetime.now(timezone.utc)
+            await log_faceswap_hit(user_email, "error", start_time, end_time)
+            raise HTTPException(500, err)
 
         with open(final_path, "rb") as f:
             result_bytes = f.read()
+
         result_key = f"bikini-theme/result/{uuid.uuid4().hex}_enhanced.png"
-        result_url = upload_to_spaces(result_bytes, result_key, content_type="image/png")
+        result_url = upload_to_spaces(result_bytes, result_key, "image/png")
 
-        #end_time = datetime.utcnow()
         end_time = datetime.now(timezone.utc)
-        await log_faceswap_hit(user_email, status="success", start_time=start_time, end_time=end_time)
+        await log_faceswap_hit(user_email, "success", start_time, end_time)
 
         return {"result_url": result_url}
 
     except Exception as e:
-        #end_time = datetime.utcnow()
         end_time = datetime.now(timezone.utc)
-
-        # Ensure we log the error with timestamps before raising
         try:
-            await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
-        except Exception as log_exc:
-            logger.error("Failed to write log_faceswap_hit: %s", log_exc)
-        raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
+            await log_faceswap_hit(user_email, "error", start_time, end_time)
+        except:
+            pass
+        raise HTTPException(500, f"Face swap failed: {e}")
+
+# @fastapi_app.post("/face-swap")
+# async def face_swap_api(
+#     source: UploadFile = File(...),
+#     target_category_id: str = Form(...),
+#     user_email: str = Depends(verify_firebase_token)
+# ):
+#     # start_time = datetime.utcnow()
+#     start_time = datetime.now(timezone.utc)
+#     try:
+#         src_bytes = await source.read()
+#         src_key = f"bikini-theme/source/{uuid.uuid4().hex}_{source.filename}"
+#         upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
+
+#         target_filename = f"{target_category_id}.png"
+#         target_url = f"https://{DO_SPACES_BUCKET}.{DO_SPACES_REGION}.digitaloceanspaces.com/bikini-theme/target/{target_filename}"
+
+#         resp = requests.get(target_url)
+#         if resp.status_code != 200:
+#             # end_time = datetime.utcnow()
+#             end_time = datetime.now(timezone.utc)
+#             await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
+#             raise HTTPException(status_code=404, detail=f"Target image not found at {target_url}")
+
+#         tgt_bytes = resp.content
+
+#         src_array = np.frombuffer(src_bytes, np.uint8)
+#         tgt_array = np.frombuffer(tgt_bytes, np.uint8)
+#         src_bgr = cv2.imdecode(src_array, cv2.IMREAD_COLOR)
+#         tgt_bgr = cv2.imdecode(tgt_array, cv2.IMREAD_COLOR)
+
+#         if src_bgr is None or tgt_bgr is None:
+#             #end_time = datetime.utcnow()
+#             end_time = datetime.now(timezone.utc)
+
+#             await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
+#             raise HTTPException(status_code=400, detail="Invalid image data")
+
+#         src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
+#         tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
+
+#         final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
+#         if err:
+#             end_time = datetime.utcnow()
+#             await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
+#             raise HTTPException(status_code=500, detail=err)
+
+#         with open(final_path, "rb") as f:
+#             result_bytes = f.read()
+#         result_key = f"bikini-theme/result/{uuid.uuid4().hex}_enhanced.png"
+#         result_url = upload_to_spaces(result_bytes, result_key, content_type="image/png")
+
+#         #end_time = datetime.utcnow()
+#         end_time = datetime.now(timezone.utc)
+#         await log_faceswap_hit(user_email, status="success", start_time=start_time, end_time=end_time)
+
+#         return {"result_url": result_url}
+
+#     except Exception as e:
+#         #end_time = datetime.utcnow()
+#         end_time = datetime.now(timezone.utc)
+
+#         # Ensure we log the error with timestamps before raising
+#         try:
+#             await log_faceswap_hit(user_email, status="error", start_time=start_time, end_time=end_time)
+#         except Exception as log_exc:
+#             logger.error("Failed to write log_faceswap_hit: %s", log_exc)
+#         raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
 
 
 @fastapi_app.get("/preview/{result_key:path}")