Update app.py

app.py

@@ -295,8 +295,8 @@ async def health():
 @fastapi_app.post("/face-swap")
 async def face_swap_api(
     source: UploadFile = File(...),
-    target_category_id: Optional[str] = Form(None), # <--- NOW OPTIONAL
-    category_id: Optional[str] = Form(None),
+    target_category_id: Optional[str] = Form(None), 
+    category_id: Optional[str] = Form(None), # <-- The ID to use for media click logging
     user_id: Optional[str] = Form(None),
     new_subcategory_id: Optional[str] = Form(None),
     user_email: str = Depends(verify_firebase_token)
@@ -306,7 +306,7 @@ async def face_swap_api(
     
     try:
         # ---------------------------------------------------------
-        # NORMALIZE EMPTY STRINGS (Android older versions)
+        # NORMALIZE EMPTY STRINGS
         # ---------------------------------------------------------
         if target_category_id == "": target_category_id = None
         if new_subcategory_id == "": new_subcategory_id = None
@@ -338,29 +338,19 @@ async def face_swap_api(
         upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
 
         # ---------------------------------------------------------
-        # TARGET IMAGE RETRIEVAL & LOGGING
+        # TARGET IMAGE RETRIEVAL
         # ---------------------------------------------------------
 
-        # CASE 1 — Old behavior (use DO Spaces target image)
         if target_provided:
+            # CASE 1 — Old behavior (use DO Spaces target image)
             target_filename = f"{target_category_id}.png"
             target_url = (
                 f"https://{DO_SPACES_BUCKET}.{DO_SPACES_REGION}."
                 f"digitaloceanspaces.com/bikini-theme/target/{target_filename}"
             )
             
-            # Log click for old system (if user_id and category_id are passed)
-            if user_id and category_id:
-                await log_media_click(user_id, category_id)
-            
-            resp = requests.get(target_url)
-            if resp.status_code != 200:
-                await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
-                raise HTTPException(status_code=404, detail=f"Target image not found: {target_url}")
-            tgt_bytes = resp.content
-
-        # CASE 2 — New behavior (use subcategory asset image)
         elif new_sub_provided:
+            # CASE 2 — New behavior (use subcategory asset image)
             
             try:
                 asset_oid = ObjectId(new_subcategory_id)
@@ -370,36 +360,38 @@ async def face_swap_api(
             # 1. Find subcategory asset by asset_images._id
             subcat_doc = await subcategories_col.find_one(
                 {"asset_images._id": asset_oid},
-                {"asset_images.$": 1, "categoryId": 1} # Need categoryId for logging
+                {"asset_images.$": 1} # Only need the asset image URL
             )
 
-            if not subcat_doc or "asset_images" not in subcat_doc or not subcat_doc.get("categoryId"):
+            if not subcat_doc or "asset_images" not in subcat_doc:
                 await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
                 raise HTTPException(
                     status_code=404,
-                    detail="Subcategory asset or parent category ID not found in DB."
+                    detail="Subcategory asset image not found in DB."
                 )
             
             asset_url = subcat_doc["asset_images"][0]["url"]
             target_url = asset_url 
-            parent_category_id = str(subcat_doc["categoryId"]) # Category ID for logging
+        
+        # ---------------------------------------------------------
+        # DOWNLOAD TARGET IMAGE
+        # ---------------------------------------------------------
+        resp = requests.get(target_url)
+        if resp.status_code != 200:
+            await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
+            raise HTTPException(status_code=404, detail=f"Target image not found or download failed: {target_url}")
+        tgt_bytes = resp.content
 
-            # Log click for new system (uses parent category ID from the subcategory doc)
-            if user_id:
-                await log_media_click(user_id, parent_category_id)
-            
-            # 2. Download target image
-            resp = requests.get(asset_url)
-            if resp.status_code != 200:
-                await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
-                raise HTTPException(
-                    status_code=404,
-                    detail=f"Failed to download asset image: {asset_url}"
-                )
-            tgt_bytes = resp.content
 
         # ---------------------------------------------------------
-        # DECODE & FACE SWAP (The rest of the logic remains the same)
+        # MEDIA CLICK LOGGING (Unified Logic)
+        # ---------------------------------------------------------
+        if user_id and category_id:
+            # Log only if both optional logging fields are provided
+            await log_media_click(user_id, category_id)
+
+        # ---------------------------------------------------------
+        # DECODE & FACE SWAP
         # ---------------------------------------------------------
         src_array = np.frombuffer(src_bytes, np.uint8)
         tgt_array = np.frombuffer(tgt_bytes, np.uint8)
@@ -414,9 +406,7 @@ async def face_swap_api(
         src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
         tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
 
-        # ---------------------------------------------------------
         # FACE SWAP & ENHANCE
-        # ---------------------------------------------------------
         final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
         if err:
             await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
@@ -451,7 +441,6 @@ async def face_swap_api(
             logger.error("Failed to write log_faceswap_hit: %s", log_exc)
             
         logger.error(f"Critical /face-swap error: {e}")
-        # The HTTPException detail might include sensitive info like target_url, simplified for the final response
         raise HTTPException(status_code=500, detail=f"Face swap failed: Internal server error.")