Update app.py

app.py

@@ -9,20 +9,16 @@ import threading
 import subprocess
 import logging
 from datetime import datetime, timezone
-
 import insightface
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
-
 from fastapi import FastAPI, UploadFile, File, HTTPException, Response, Depends, Security, Form
 from fastapi.responses import RedirectResponse
 from pydantic import BaseModel
 from motor.motor_asyncio import AsyncIOMotorClient
-
 import uvicorn
 import gradio as gr
 from gradio import mount_gradio_app
-
 # DigitalOcean Spaces
 import boto3
 from botocore.client import Config
@@ -31,40 +27,32 @@ from typing import Optional
 import requests
 import json
 from bson import ObjectId  
-
 # --------------------- Logging ---------------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-
 # --------------------- Paths -----------------------
 REPO_ID = "HariLogicgo/face_swap_models"
 BASE_DIR = "./workspace"
 MODELS_DIR = "./models"
-
 os.makedirs(MODELS_DIR, exist_ok=True)
-
 # --------------------- Secrets ---------------------
 HF_TOKEN = os.getenv("HF_TOKEN")  # Hugging Face private repo token
 # Firebase credentials JSON
 FIREBASE_CREDENTIALS_PATH = os.getenv("FIREBASE_CREDENTIALS_PATH")
-
 # --------------------- DigitalOcean Spaces Credentials ---------------------
 DO_SPACES_REGION = os.getenv("DO_SPACES_REGION", "blr1")
 DO_SPACES_ENDPOINT = os.getenv("DO_SPACES_ENDPOINT", f"https://{DO_SPACES_REGION}.digitaloceanspaces.com")
 DO_SPACES_KEY = os.getenv("DO_SPACES_KEY")
 DO_SPACES_SECRET = os.getenv("DO_SPACES_SECRET")
 DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET")
-
 # --------------------- Firebase Auth ---------------------
 import firebase_admin
 from firebase_admin import credentials, auth
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-
 if not firebase_admin._apps:
     FIREBASE_CREDENTIALS = os.getenv("FIREBASE_CREDENTIALS_PATH")
     if not FIREBASE_CREDENTIALS:
         raise RuntimeError("❌ FIREBASE_CREDENTIALS_PATH not set in environment variables")
-
     try:
         # Try parsing as JSON string
         cred_dict = json.loads(FIREBASE_CREDENTIALS)
@@ -74,11 +62,8 @@ if not firebase_admin._apps:
         # Fallback: assume it's a file path
         cred = credentials.Certificate(FIREBASE_CREDENTIALS)
         logger.info("✅ Firebase initialized from JSON file path")
-
     firebase_admin.initialize_app(cred)
-
 security = HTTPBearer()
-
 def verify_firebase_token(credentials: HTTPAuthorizationCredentials = Security(security)):
     """Verify Firebase ID token from Authorization header."""
     try:
@@ -91,7 +76,6 @@ def verify_firebase_token(credentials: HTTPAuthorizationCredentials = Security(s
     except Exception as e:
         logger.error(f"Firebase auth failed: {e}")
         raise HTTPException(status_code=401, detail="Unauthorized: Invalid Firebase token")
-
 # --------------------- Download Models ---------------------
 def download_models():
     logger.info("Downloading models from private HF repo...")
@@ -102,7 +86,6 @@ def download_models():
         local_dir=MODELS_DIR,
         token=HF_TOKEN
     )
-
     buffalo_files = [
         "1k3d68.onnx",
         "2d106det.onnx",
@@ -120,18 +103,14 @@ def download_models():
         )
     logger.info("Models downloaded successfully")
     return inswapper_path
-
 inswapper_path = download_models()
-
 # --------------------- Face Analysis + Swapper ---------------------
 providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
 face_analysis_app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=providers)
 face_analysis_app.prepare(ctx_id=0, det_size=(640, 640))
 swapper = insightface.model_zoo.get_model(inswapper_path, providers=providers)
-
 # --------------------- CodeFormer ---------------------
 CODEFORMER_PATH = "CodeFormer/inference_codeformer.py"
-
 def ensure_codeformer():
     if not os.path.exists("CodeFormer"):
         subprocess.run("git clone https://github.com/sczhou/CodeFormer.git", shell=True, check=True)
@@ -139,30 +118,21 @@ def ensure_codeformer():
         subprocess.run("python CodeFormer/basicsr/setup.py develop", shell=True, check=True)
         subprocess.run("python CodeFormer/scripts/download_pretrained_models.py facelib", shell=True, check=True)
         subprocess.run("python CodeFormer/scripts/download_pretrained_models.py CodeFormer", shell=True, check=True)
-
 ensure_codeformer()
-
 # --------------------- MongoDB ---------------------
 MONGODB_URL = os.getenv("MONGODB_URL")
-
 client = None
 database = None
-
-# --------------------- Admin Panel DB (categories + media_clicks) ---------------------
 # --------------------- Admin Panel DB (categories + subcategories + media_clicks) ---------------------
 ADMIN_MONGO_URL = os.getenv("ADMIN_MONGO_URL")
 admin_client = AsyncIOMotorClient(ADMIN_MONGO_URL)
 admin_db = admin_client.adminPanel
-
 # Collections
 categories_col = admin_db.categories
 subcategories_col = admin_db.subcategories
 media_clicks_col = admin_db.media_clicks
-users_col = admin_db.users  # optional, only if needed
-
 # --------------------- FastAPI ---------------------
 fastapi_app = FastAPI()
-
 @fastapi_app.on_event("startup")
 async def startup_db():
     global client, database
@@ -170,22 +140,18 @@ async def startup_db():
     client = AsyncIOMotorClient(MONGODB_URL)
     database = client.FaceSwap
     logger.info("MongoDB initialized for API logs")
-
 @fastapi_app.on_event("shutdown")
 async def shutdown_db():
     global client
     if client:
         client.close()
         logger.info("MongoDB connection closed")
-
 # --------------------- Logging API Hits ---------------------
 async def log_faceswap_hit(user_email: str, status: str, start_time: datetime, end_time: datetime):
     global database
     if database is None:
         return
-
     response_time_ms = (end_time - start_time).total_seconds() * 1000
-
     await database.api_logs.insert_one({
         "user": user_email,
         "endpoint": "/face-swap",
@@ -195,73 +161,109 @@ async def log_faceswap_hit(user_email: str, status: str, start_time: datetime, e
         "response_time_ms": response_time_ms
     })
 
+# --------------------- Media Click Logging Helper ---------------------
+async def log_media_click(user_id: str, category_oid_str: str):
+    """
+    Logs a click event to the media_clicks collection against the Category.
+    """
+    try:
+        user_oid = ObjectId(user_id.strip())
+        category_oid = ObjectId(category_oid_str.strip())
+        now = datetime.utcnow()
+
+        # Check if category exists (optional, but prevents logging invalid IDs)
+        if not await categories_col.find_one({"_id": category_oid}):
+            logger.warning(f"Category ID {category_oid_str} not found. Skipping click logging.")
+            return
+
+        # 1) Try updating existing category click entry
+        update_result = await media_clicks_col.update_one(
+            {
+                "userId": user_oid,
+                "categories.categoryId": category_oid
+            },
+            {
+                "$set": {
+                    "updatedAt": now,
+                    "categories.$.lastClickedAt": now
+                },
+                "$inc": {
+                    "categories.$.click_count": 1
+                }
+            }
+        )
+
+        # 2) If no category entry exists → push new one (or create doc)
+        if update_result.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
+            )
+        
+        logger.info(f"Media click logged for User {user_id} on Category {category_oid_str}")
+
+    except Exception as media_err:
+        logger.error(f"MEDIA_CLICK LOGGING ERROR: {media_err}")
 
 # --------------------- Face Swap Pipeline ---------------------
 swap_lock = threading.Lock()
-
 def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
     try:
         with swap_lock:
             if os.path.exists(temp_dir):
                 shutil.rmtree(temp_dir)
             os.makedirs(temp_dir, exist_ok=True)
-
             src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
             tgt_bgr_full = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
-
             src_faces = face_analysis_app.get(src_bgr)
             tgt_faces = face_analysis_app.get(tgt_bgr_full)
-
             if not src_faces or not tgt_faces:
                 return None, None, "❌ Face not detected in source or target image"
-
             src_face0 = src_faces[0]
             tgt_face0 = tgt_faces[0]
-
             swapped_bgr_full = swapper.get(tgt_bgr_full, tgt_face0, src_face0)
             if swapped_bgr_full is None:
                 return None, None, "❌ Face swap failed"
-
             swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
             cv2.imwrite(swapped_path, swapped_bgr_full)
-
             cmd = f"python {CODEFORMER_PATH} -w 0.7 --input_path {swapped_path} --output_path {temp_dir} --bg_upsampler realesrgan --face_upsample"
             result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
             if result.returncode != 0:
                 return None, None, f"❌ CodeFormer failed:\n{result.stderr}"
-
             final_results_dir = os.path.join(temp_dir, "final_results")
             final_files = [f for f in os.listdir(final_results_dir) if f.endswith(".png")]
             if not final_files:
                 return None, None, "❌ No enhanced image found"
-
             final_path = os.path.join(final_results_dir, final_files[0])
             final_img = cv2.cvtColor(cv2.imread(final_path), cv2.COLOR_BGR2RGB)
-
             return final_img, final_path, ""
-
     except Exception as e:
         return None, None, f"❌ Error: {str(e)}"
-
 # --------------------- Gradio ---------------------
 with gr.Blocks() as demo:
     gr.Markdown("Face Swap")
-
     with gr.Row():
         src_input = gr.Image(type="numpy", label="Upload Your Face")
         tgt_input = gr.Image(type="numpy", label="Upload Target Image")
-
     btn = gr.Button("Swap Face")
     output_img = gr.Image(type="numpy", label="Enhanced Output")
     download = gr.File(label="⬇️ Download Enhanced Image")
     error_box = gr.Textbox(label="Logs / Errors", interactive=False)
-
     def process(src, tgt):
         img, path, err = face_swap_and_enhance(src, tgt)
         return img, path, err
-
     btn.click(process, [src_input, tgt_input], [output_img, download, error_box])
-
 # --------------------- DigitalOcean Spaces Helper ---------------------
 def get_spaces_client():
     session = boto3.session.Session()
@@ -274,22 +276,18 @@ def get_spaces_client():
         config=Config(signature_version='s3v4')
     )
     return client
-
 def upload_to_spaces(file_bytes, key, content_type="image/png"):
     client = get_spaces_client()
     client.put_object(Bucket=DO_SPACES_BUCKET, Key=key, Body=file_bytes, ContentType=content_type, ACL='public-read')
     return f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{key}"
-
 def download_from_spaces(key):
     client = get_spaces_client()
     obj = client.get_object(Bucket=DO_SPACES_BUCKET, Key=key)
     return obj['Body'].read()
-
 # --------------------- API Endpoints ---------------------
 @fastapi_app.get("/")
 def root():
     return RedirectResponse("/gradio")
-
 @fastapi_app.get("/health")
 async def health():
     return {"status": "healthy"}
@@ -297,40 +295,36 @@ async def health():
 @fastapi_app.post("/face-swap")
 async def face_swap_api(
     source: UploadFile = File(...),
-    target_category_id: str = Form(...),  # REQUIRED (old behavior preserved)
+    target_category_id: Optional[str] = Form(None), # <--- NOW OPTIONAL
     category_id: Optional[str] = Form(None),
     user_id: Optional[str] = Form(None),
     new_subcategory_id: Optional[str] = Form(None),
     user_email: str = Depends(verify_firebase_token)
 ):
     start_time = datetime.now(timezone.utc)
-
+    target_url = None
+    
     try:
         # ---------------------------------------------------------
         # NORMALIZE EMPTY STRINGS (Android older versions)
         # ---------------------------------------------------------
-        if target_category_id == "":
-            target_category_id = None
-
-        if new_subcategory_id == "":
-            new_subcategory_id = None
-
-        if category_id == "":
-            category_id = None
-
-        if user_id == "":
-            user_id = None
-
+        if target_category_id == "": target_category_id = None
+        if new_subcategory_id == "": new_subcategory_id = None
+        if category_id == "": category_id = None
+        if user_id == "": user_id = None
+        
         # ---------------------------------------------------------
         # STRICT XOR VALIDATION
         # ---------------------------------------------------------
-        if target_category_id and new_subcategory_id:
+        target_provided = target_category_id is not None
+        new_sub_provided = new_subcategory_id is not None
+
+        if target_provided and new_sub_provided:
             raise HTTPException(
                 status_code=400,
                 detail="Provide ONLY ONE of: target_category_id OR new_subcategory_id"
             )
-
-        if not target_category_id and not new_subcategory_id:
+        if not target_provided and not new_sub_provided:
             raise HTTPException(
                 status_code=400,
                 detail="Either target_category_id OR new_subcategory_id is required"
@@ -344,57 +338,71 @@ async def face_swap_api(
         upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
 
         # ---------------------------------------------------------
-        # CASE 1 — Old behavior (use DO Spaces target image)
+        # TARGET IMAGE RETRIEVAL & LOGGING
         # ---------------------------------------------------------
-        if target_category_id:
+
+        # CASE 1 — Old behavior (use DO Spaces target image)
+        if target_provided:
             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)
-        # ---------------------------------------------------------
-        else:
-            # Find subcategory asset by asset_images._id
-            asset = await admin_db.subcategories.find_one(
-                {"asset_images._id": ObjectId(new_subcategory_id)},
-                {"asset_images.$": 1}
+        elif new_sub_provided:
+            
+            try:
+                asset_oid = ObjectId(new_subcategory_id)
+            except:
+                raise HTTPException(400, "Invalid new_subcategory_id format.")
+
+            # 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
             )
 
-            if not asset or "asset_images" not in asset:
+            if not subcat_doc or "asset_images" not in subcat_doc or not subcat_doc.get("categoryId"):
                 await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
                 raise HTTPException(
                     status_code=404,
-                    detail="Subcategory asset image not found"
+                    detail="Subcategory asset or parent category ID not found in DB."
                 )
-
-            # Extract the single matching image URL
-            asset_url = asset["asset_images"][0]["url"]
-
+            
+            asset_url = subcat_doc["asset_images"][0]["url"]
+            target_url = asset_url 
+            parent_category_id = str(subcat_doc["categoryId"]) # Category ID for logging
+
+            # 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 BOTH IMAGES
+        # DECODE & FACE SWAP (The rest of the logic remains the same)
         # ---------------------------------------------------------
         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)
 
@@ -402,6 +410,7 @@ async def face_swap_api(
             await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
             raise HTTPException(status_code=400, detail="Invalid image data")
 
+        # Convert to RGB for processing
         src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
         tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
 
@@ -416,111 +425,584 @@ async def face_swap_api(
         # Save final output to DO Spaces
         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, "image/png")
-
-        await log_faceswap_hit(user_email, "success", start_time, datetime.now(timezone.utc))
-
+        
         # ---------------------------------------------------------
         # SUCCESS RESPONSE
         # ---------------------------------------------------------
+        await log_faceswap_hit(user_email, "success", start_time, datetime.now(timezone.utc))
+        
         return {
             "result_url": result_url,
-            "category_id": category_id,
+            "category_id": category_id, 
             "user_id": user_id,
             "new_subcategory_id": new_subcategory_id
         }
 
+    except HTTPException:
+        raise
+        
     except Exception as e:
-        await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
-        raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
+        end_time = datetime.now(timezone.utc)
+        try:
+            await log_faceswap_hit(user_email, "error", start_time, end_time)
+        except Exception as log_exc:
+            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.")
 
 
-####------------------------------------OLD CODE------------------------------------####
+@fastapi_app.get("/preview/{result_key:path}")
+async def preview_result(result_key: str):
+    try:
+        img_bytes = download_from_spaces(result_key)
+    except Exception:
+        raise HTTPException(status_code=404, detail="Result not found")
+    return Response(
+        content=img_bytes,
+        media_type="image/png",
+        headers={"Content-Disposition": "inline; filename=result.png"}
+    )
+# --------------------- Mount Gradio ---------------------
+fastapi_app = mount_gradio_app(fastapi_app, demo, path="/gradio")
+if __name__ == "__main__":
+    uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)
+
+
+
+
+
+
+
+
+
+
+# # --------------------- List Images Endpoint ---------------------
+# import os
+# os.environ["OMP_NUM_THREADS"] = "1"
+# import shutil
+# import uuid
+# import cv2
+# import numpy as np
+# import threading
+# import subprocess
+# import logging
+# from datetime import datetime, timezone
+
+# import insightface
+# from insightface.app import FaceAnalysis
+# from huggingface_hub import hf_hub_download
+
+# from fastapi import FastAPI, UploadFile, File, HTTPException, Response, Depends, Security, Form
+# from fastapi.responses import RedirectResponse
+# from pydantic import BaseModel
+# from motor.motor_asyncio import AsyncIOMotorClient
+
+# import uvicorn
+# import gradio as gr
+# from gradio import mount_gradio_app
+
+# # DigitalOcean Spaces
+# import boto3
+# from botocore.client import Config
+# from io import BytesIO
+# from typing import Optional
+# import requests
+# import json
+# from bson import ObjectId  
+
+# # --------------------- Logging ---------------------
+# logging.basicConfig(level=logging.INFO)
+# logger = logging.getLogger(__name__)
+
+# # --------------------- Paths -----------------------
+# REPO_ID = "HariLogicgo/face_swap_models"
+# BASE_DIR = "./workspace"
+# MODELS_DIR = "./models"
+
+# os.makedirs(MODELS_DIR, exist_ok=True)
+
+# # --------------------- Secrets ---------------------
+# HF_TOKEN = os.getenv("HF_TOKEN")  # Hugging Face private repo token
+# # Firebase credentials JSON
+# FIREBASE_CREDENTIALS_PATH = os.getenv("FIREBASE_CREDENTIALS_PATH")
+
+# # --------------------- DigitalOcean Spaces Credentials ---------------------
+# DO_SPACES_REGION = os.getenv("DO_SPACES_REGION", "blr1")
+# DO_SPACES_ENDPOINT = os.getenv("DO_SPACES_ENDPOINT", f"https://{DO_SPACES_REGION}.digitaloceanspaces.com")
+# DO_SPACES_KEY = os.getenv("DO_SPACES_KEY")
+# DO_SPACES_SECRET = os.getenv("DO_SPACES_SECRET")
+# DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET")
+
+# # --------------------- Firebase Auth ---------------------
+# import firebase_admin
+# from firebase_admin import credentials, auth
+# from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+
+# if not firebase_admin._apps:
+#     FIREBASE_CREDENTIALS = os.getenv("FIREBASE_CREDENTIALS_PATH")
+#     if not FIREBASE_CREDENTIALS:
+#         raise RuntimeError("❌ FIREBASE_CREDENTIALS_PATH not set in environment variables")
+
+#     try:
+#         # Try parsing as JSON string
+#         cred_dict = json.loads(FIREBASE_CREDENTIALS)
+#         cred = credentials.Certificate(cred_dict)
+#         logger.info("✅ Firebase initialized from JSON string in environment variable")
+#     except json.JSONDecodeError:
+#         # Fallback: assume it's a file path
+#         cred = credentials.Certificate(FIREBASE_CREDENTIALS)
+#         logger.info("✅ Firebase initialized from JSON file path")
+
+#     firebase_admin.initialize_app(cred)
+
+# security = HTTPBearer()
+
+# def verify_firebase_token(credentials: HTTPAuthorizationCredentials = Security(security)):
+#     """Verify Firebase ID token from Authorization header."""
+#     try:
+#         id_token = credentials.credentials
+#         decoded_token = auth.verify_id_token(id_token)
+#         user_email = decoded_token.get("email")
+#         if not user_email:
+#             raise HTTPException(status_code=401, detail="Firebase token invalid or missing email")
+#         return user_email
+#     except Exception as e:
+#         logger.error(f"Firebase auth failed: {e}")
+#         raise HTTPException(status_code=401, detail="Unauthorized: Invalid Firebase token")
+
+# # --------------------- Download Models ---------------------
+# def download_models():
+#     logger.info("Downloading models from private HF repo...")
+#     inswapper_path = hf_hub_download(
+#         repo_id=REPO_ID,
+#         filename="models/inswapper_128.onnx",
+#         repo_type="model",
+#         local_dir=MODELS_DIR,
+#         token=HF_TOKEN
+#     )
+
+#     buffalo_files = [
+#         "1k3d68.onnx",
+#         "2d106det.onnx",
+#         "genderage.onnx",
+#         "det_10g.onnx",
+#         "w600k_r50.onnx"
+#     ]
+#     for f in buffalo_files:
+#         hf_hub_download(
+#             repo_id=REPO_ID,
+#             filename=f"models/buffalo_l/{f}",
+#             repo_type="model",
+#             local_dir=MODELS_DIR,
+#             token=HF_TOKEN
+#         )
+#     logger.info("Models downloaded successfully")
+#     return inswapper_path
+
+# inswapper_path = download_models()
+
+# # --------------------- Face Analysis + Swapper ---------------------
+# providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+# face_analysis_app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=providers)
+# face_analysis_app.prepare(ctx_id=0, det_size=(640, 640))
+# swapper = insightface.model_zoo.get_model(inswapper_path, providers=providers)
+
+# # --------------------- CodeFormer ---------------------
+# CODEFORMER_PATH = "CodeFormer/inference_codeformer.py"
+
+# def ensure_codeformer():
+#     if not os.path.exists("CodeFormer"):
+#         subprocess.run("git clone https://github.com/sczhou/CodeFormer.git", shell=True, check=True)
+#         subprocess.run("pip install -r CodeFormer/requirements.txt", shell=True, check=True)
+#         subprocess.run("python CodeFormer/basicsr/setup.py develop", shell=True, check=True)
+#         subprocess.run("python CodeFormer/scripts/download_pretrained_models.py facelib", shell=True, check=True)
+#         subprocess.run("python CodeFormer/scripts/download_pretrained_models.py CodeFormer", shell=True, check=True)
+
+# ensure_codeformer()
+
+# # --------------------- MongoDB ---------------------
+# MONGODB_URL = os.getenv("MONGODB_URL")
+
+# client = None
+# database = None
+
+# # --------------------- Admin Panel DB (categories + media_clicks) ---------------------
+# # --------------------- Admin Panel DB (categories + subcategories + media_clicks) ---------------------
+# ADMIN_MONGO_URL = os.getenv("ADMIN_MONGO_URL")
+# admin_client = AsyncIOMotorClient(ADMIN_MONGO_URL)
+# admin_db = admin_client.adminPanel
+
+# # Collections
+# categories_col = admin_db.categories
+# subcategories_col = admin_db.subcategories
+# media_clicks_col = admin_db.media_clicks
+# users_col = admin_db.users  # optional, only if needed
+
+# # --------------------- FastAPI ---------------------
+# fastapi_app = FastAPI()
+
+# @fastapi_app.on_event("startup")
+# async def startup_db():
+#     global client, database
+#     logger.info("Initializing MongoDB for API logs...")
+#     client = AsyncIOMotorClient(MONGODB_URL)
+#     database = client.FaceSwap
+#     logger.info("MongoDB initialized for API logs")
+
+# @fastapi_app.on_event("shutdown")
+# async def shutdown_db():
+#     global client
+#     if client:
+#         client.close()
+#         logger.info("MongoDB connection closed")
+
+# # --------------------- Logging API Hits ---------------------
+# async def log_faceswap_hit(user_email: str, status: str, start_time: datetime, end_time: datetime):
+#     global database
+#     if database is None:
+#         return
+
+#     response_time_ms = (end_time - start_time).total_seconds() * 1000
+
+#     await database.api_logs.insert_one({
+#         "user": user_email,
+#         "endpoint": "/face-swap",
+#         "status": status,
+#         "start_time": start_time,
+#         "end_time": end_time,
+#         "response_time_ms": response_time_ms
+#     })
+
+
+# # --------------------- Face Swap Pipeline ---------------------
+# swap_lock = threading.Lock()
+
+# def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
+#     try:
+#         with swap_lock:
+#             if os.path.exists(temp_dir):
+#                 shutil.rmtree(temp_dir)
+#             os.makedirs(temp_dir, exist_ok=True)
+
+#             src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
+#             tgt_bgr_full = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
+
+#             src_faces = face_analysis_app.get(src_bgr)
+#             tgt_faces = face_analysis_app.get(tgt_bgr_full)
+
+#             if not src_faces or not tgt_faces:
+#                 return None, None, "❌ Face not detected in source or target image"
+
+#             src_face0 = src_faces[0]
+#             tgt_face0 = tgt_faces[0]
+
+#             swapped_bgr_full = swapper.get(tgt_bgr_full, tgt_face0, src_face0)
+#             if swapped_bgr_full is None:
+#                 return None, None, "❌ Face swap failed"
+
+#             swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
+#             cv2.imwrite(swapped_path, swapped_bgr_full)
+
+#             cmd = f"python {CODEFORMER_PATH} -w 0.7 --input_path {swapped_path} --output_path {temp_dir} --bg_upsampler realesrgan --face_upsample"
+#             result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+#             if result.returncode != 0:
+#                 return None, None, f"❌ CodeFormer failed:\n{result.stderr}"
+
+#             final_results_dir = os.path.join(temp_dir, "final_results")
+#             final_files = [f for f in os.listdir(final_results_dir) if f.endswith(".png")]
+#             if not final_files:
+#                 return None, None, "❌ No enhanced image found"
+
+#             final_path = os.path.join(final_results_dir, final_files[0])
+#             final_img = cv2.cvtColor(cv2.imread(final_path), cv2.COLOR_BGR2RGB)
+
+#             return final_img, final_path, ""
+
+#     except Exception as e:
+#         return None, None, f"❌ Error: {str(e)}"
+
+# # --------------------- Gradio ---------------------
+# with gr.Blocks() as demo:
+#     gr.Markdown("Face Swap")
+
+#     with gr.Row():
+#         src_input = gr.Image(type="numpy", label="Upload Your Face")
+#         tgt_input = gr.Image(type="numpy", label="Upload Target Image")
+
+#     btn = gr.Button("Swap Face")
+#     output_img = gr.Image(type="numpy", label="Enhanced Output")
+#     download = gr.File(label="⬇️ Download Enhanced Image")
+#     error_box = gr.Textbox(label="Logs / Errors", interactive=False)
+
+#     def process(src, tgt):
+#         img, path, err = face_swap_and_enhance(src, tgt)
+#         return img, path, err
+
+#     btn.click(process, [src_input, tgt_input], [output_img, download, error_box])
+
+# # --------------------- DigitalOcean Spaces Helper ---------------------
+# def get_spaces_client():
+#     session = boto3.session.Session()
+#     client = session.client(
+#         's3',
+#         region_name=DO_SPACES_REGION,
+#         endpoint_url=DO_SPACES_ENDPOINT,
+#         aws_access_key_id=DO_SPACES_KEY,
+#         aws_secret_access_key=DO_SPACES_SECRET,
+#         config=Config(signature_version='s3v4')
+#     )
+#     return client
+
+# def upload_to_spaces(file_bytes, key, content_type="image/png"):
+#     client = get_spaces_client()
+#     client.put_object(Bucket=DO_SPACES_BUCKET, Key=key, Body=file_bytes, ContentType=content_type, ACL='public-read')
+#     return f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{key}"
+
+# def download_from_spaces(key):
+#     client = get_spaces_client()
+#     obj = client.get_object(Bucket=DO_SPACES_BUCKET, Key=key)
+#     return obj['Body'].read()
+
+# # --------------------- API Endpoints ---------------------
+# @fastapi_app.get("/")
+# def root():
+#     return RedirectResponse("/gradio")
+
+# @fastapi_app.get("/health")
+# async def health():
+#     return {"status": "healthy"}
+
 # @fastapi_app.post("/face-swap")
 # async def face_swap_api(
 #     source: UploadFile = File(...),
-#     target_category_id: str = Form(...),
+#     target_category_id: str = Form(...),  # REQUIRED (old behavior preserved)
+#     category_id: Optional[str] = Form(None),
+#     user_id: Optional[str] = Form(None),
+#     new_subcategory_id: Optional[str] = Form(None),
 #     user_email: str = Depends(verify_firebase_token)
 # ):
-#     # start_time = datetime.utcnow()
 #     start_time = datetime.now(timezone.utc)
+
 #     try:
+#         # ---------------------------------------------------------
+#         # NORMALIZE EMPTY STRINGS (Android older versions)
+#         # ---------------------------------------------------------
+#         if target_category_id == "":
+#             target_category_id = None
+
+#         if new_subcategory_id == "":
+#             new_subcategory_id = None
+
+#         if category_id == "":
+#             category_id = None
+
+#         if user_id == "":
+#             user_id = None
+
+#         # ---------------------------------------------------------
+#         # STRICT XOR VALIDATION
+#         # ---------------------------------------------------------
+#         if target_category_id and new_subcategory_id:
+#             raise HTTPException(
+#                 status_code=400,
+#                 detail="Provide ONLY ONE of: target_category_id OR new_subcategory_id"
+#             )
+
+#         if not target_category_id and not new_subcategory_id:
+#             raise HTTPException(
+#                 status_code=400,
+#                 detail="Either target_category_id OR new_subcategory_id is required"
+#             )
+
+#         # ---------------------------------------------------------
+#         # READ SOURCE IMAGE
+#         # ---------------------------------------------------------
 #         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
-
+#         # ---------------------------------------------------------
+#         # CASE 1 — Old behavior (use DO Spaces target image)
+#         # ---------------------------------------------------------
+#         if target_category_id:
+#             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}"
+#             )
+
+#             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)
+#         # ---------------------------------------------------------
+#         else:
+#             # Find subcategory asset by asset_images._id
+#             asset = await admin_db.subcategories.find_one(
+#                 {"asset_images._id": ObjectId(new_subcategory_id)},
+#                 {"asset_images.$": 1}
+#             )
+
+#             if not asset or "asset_images" not in asset:
+#                 await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
+#                 raise HTTPException(
+#                     status_code=404,
+#                     detail="Subcategory asset image not found"
+#                 )
+
+#             # Extract the single matching image URL
+#             asset_url = asset["asset_images"][0]["url"]
+
+#             resp = requests.get(asset_url)
+#             if resp.status_code != 200:
+#                 raise HTTPException(
+#                     status_code=404,
+#                     detail=f"Failed to download asset image: {asset_url}"
+#                 )
+
+#             tgt_bytes = resp.content
+
+#         # ---------------------------------------------------------
+#         # DECODE BOTH IMAGES
+#         # ---------------------------------------------------------
 #         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)
+#             await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
 #             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)
 
+#         # ---------------------------------------------------------
+#         # FACE SWAP & ENHANCE
+#         # ---------------------------------------------------------
 #         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)
+#             await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
 #             raise HTTPException(status_code=500, detail=err)
 
+#         # Save final output to DO Spaces
 #         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, datetime.now(timezone.utc))
 
-#         return {"result_url": result_url}
+#         # ---------------------------------------------------------
+#         # SUCCESS RESPONSE
+#         # ---------------------------------------------------------
+#         return {
+#             "result_url": result_url,
+#             "category_id": category_id,
+#             "user_id": user_id,
+#             "new_subcategory_id": new_subcategory_id
+#         }
 
 #     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)
+#         await log_faceswap_hit(user_email, "error", start_time, datetime.now(timezone.utc))
 #         raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
 
 
-@fastapi_app.get("/preview/{result_key:path}")
-async def preview_result(result_key: str):
-    try:
-        img_bytes = download_from_spaces(result_key)
-    except Exception:
-        raise HTTPException(status_code=404, detail="Result not found")
-    return Response(
-        content=img_bytes,
-        media_type="image/png",
-        headers={"Content-Disposition": "inline; filename=result.png"}
-    )
-
-# --------------------- Mount Gradio ---------------------
-fastapi_app = mount_gradio_app(fastapi_app, demo, path="/gradio")
-
-if __name__ == "__main__":
-    uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)
+# ####------------------------------------OLD CODE------------------------------------####
+# # @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}")
+# async def preview_result(result_key: str):
+#     try:
+#         img_bytes = download_from_spaces(result_key)
+#     except Exception:
+#         raise HTTPException(status_code=404, detail="Result not found")
+#     return Response(
+#         content=img_bytes,
+#         media_type="image/png",
+#         headers={"Content-Disposition": "inline; filename=result.png"}
+#     )
+
+# # --------------------- Mount Gradio ---------------------
+# fastapi_app = mount_gradio_app(fastapi_app, demo, path="/gradio")
+
+# if __name__ == "__main__":
+#     uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)