Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import io
 import json
 import traceback
-from datetime import datetime
+from datetime import datetime,timedelta
 from typing import Optional
 import time
 from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request, Depends
@@ -173,12 +173,8 @@ async def generate(
         combined_img = pil_img1
 
     # -------------------------
-    # 3. OPTIONAL CATEGORY CLICK LOGIC
+    # 3. CATEGORY CLICK LOGIC
     # -------------------------
-    # -------------------------
-    # 3. OPTIONAL CATEGORY CLICK LOGIC
-    # -------------------------
-    
     if user_id and category_id:
         try:
             admin_client = MongoClient(os.getenv("ADMIN_MONGODB_URI"))
@@ -198,26 +194,58 @@ async def generate(
     
             now = datetime.utcnow()
     
+            # Normalize dates (UTC midnight)
+            today_date = datetime(now.year, now.month, now.day)
+            yesterday_date = today_date - timedelta(days=1)
+    
             # --------------------------------------------------
             # AI EDIT USAGE TRACKING (GLOBAL PER USER)
             # --------------------------------------------------
             media_clicks_col.update_one(
                 {"userId": user_oid},
                 {
+                    "$setOnInsert": {
+                        "createdAt": now,
+                        "ai_edit_daily_count": []
+                    },
                     "$set": {
                         "ai_edit_last_date": now,
                         "updatedAt": now
                     },
                     "$inc": {
                         "ai_edit_complete": 1
-                    },
-                    "$setOnInsert": {
-                        "createdAt": now
                     }
                 },
                 upsert=True
             )
     
+            # --------------------------------------------------
+            # DAILY COUNT LOGIC
+            # --------------------------------------------------
+            doc = media_clicks_col.find_one({"userId": user_oid}, {"ai_edit_daily_count": 1})
+            daily_entries = doc.get("ai_edit_daily_count", []) if doc else []
+    
+            is_first_time = len(daily_entries) == 0
+            existing_dates = {entry["date"].date(): entry["count"] for entry in daily_entries}
+            daily_updates = []
+    
+            if is_first_time:
+                # First-ever usage → only today
+                daily_updates.append({"date": today_date, "count": 1})
+            else:
+                # Backfill yesterday if missing
+                if yesterday_date.date() not in existing_dates:
+                    daily_updates.append({"date": yesterday_date, "count": 0})
+                # Insert today if missing
+                if today_date.date() not in existing_dates:
+                    daily_updates.append({"date": today_date, "count": 1})
+    
+            if daily_updates:
+                media_clicks_col.update_one(
+                    {"userId": user_oid},
+                    {"$push": {"ai_edit_daily_count": {"$each": daily_updates}}}
+                )
+    
             # --------------------------------------------------
             # CATEGORY CLICK LOGIC
             # --------------------------------------------------
@@ -253,29 +281,49 @@ async def generate(
     
         except Exception as e:
             print("CATEGORY_LOG_ERROR:", e)
-            # DO NOT break image generation
-
     
     # if user_id and category_id:
     #     try:
     #         admin_client = MongoClient(os.getenv("ADMIN_MONGODB_URI"))
     #         admin_db = admin_client["adminPanel"]
-
+    
     #         categories_col = admin_db.categories
     #         media_clicks_col = admin_db.media_clicks
-
+    
     #         # Validate user_oid & category_oid
     #         user_oid = ObjectId(user_id)
     #         category_oid = ObjectId(category_id)
-
+    
     #         # Check category exists
     #         category_doc = categories_col.find_one({"_id": category_oid})
     #         if not category_doc:
     #             raise HTTPException(400, f"Invalid category_id: {category_id}")
-
+    
     #         now = datetime.utcnow()
-
-    #         # Try updating existing category click
+    
+    #         # --------------------------------------------------
+    #         # AI EDIT USAGE TRACKING (GLOBAL PER USER)
+    #         # --------------------------------------------------
+    #         media_clicks_col.update_one(
+    #             {"userId": user_oid},
+    #             {
+    #                 "$set": {
+    #                     "ai_edit_last_date": now,
+    #                     "updatedAt": now
+    #                 },
+    #                 "$inc": {
+    #                     "ai_edit_complete": 1
+    #                 },
+    #                 "$setOnInsert": {
+    #                     "createdAt": now
+    #                 }
+    #             },
+    #             upsert=True
+    #         )
+    
+    #         # --------------------------------------------------
+    #         # CATEGORY CLICK LOGIC
+    #         # --------------------------------------------------
     #         update_res = media_clicks_col.update_one(
     #             {"userId": user_oid, "categories.categoryId": category_oid},
     #             {
@@ -283,16 +331,17 @@ async def generate(
     #                     "updatedAt": now,
     #                     "categories.$.lastClickedAt": now
     #                 },
-    #                 "$inc": {"categories.$.click_count": 1}
+    #                 "$inc": {
+    #                     "categories.$.click_count": 1
+    #                 }
     #             }
     #         )
-
-    #         # If user or category does not exist → create/push new
+    
+    #         # If category does not exist → push new
     #         if update_res.matched_count == 0:
     #             media_clicks_col.update_one(
     #                 {"userId": user_oid},
     #                 {
-    #                     "$setOnInsert": {"createdAt": now},
     #                     "$set": {"updatedAt": now},
     #                     "$push": {
     #                         "categories": {
@@ -304,11 +353,9 @@ async def generate(
     #                 },
     #                 upsert=True
     #             )
-
+    
     #     except Exception as e:
     #         print("CATEGORY_LOG_ERROR:", e)
-            # but do NOT stop the image generation
-
     # -------------------------
     # 4. HF INFERENCE
     # -------------------------