import os from datetime import datetime, timedelta from typing import Optional, Union from urllib.parse import urlparse from pymongo import MongoClient, ReturnDocument from pymongo.errors import DuplicateKeyError from bson import ObjectId import config # Global dictionary to cache MongoDB clients mongodb_clients = {} def _normalize_date(dt: datetime) -> datetime: """Return UTC date at midnight for consistent daily counters.""" return datetime(dt.year, dt.month, dt.day) def _build_ai_edit_daily_count(existing: Optional[list], now: datetime) -> list: """ Ensure ai_edit_daily_count is present and includes all days up to today. - If today's date exists: return list unchanged. - If missing dates: fill gaps with count=0 and add today with count=1. - Always keep only the most recent 32 entries (oldest removed first). """ today = _normalize_date(now) if not existing: return [{"date": today, "count": 1}] # Normalize dates and keep order normalized = [] seen_dates = set() for entry in existing: date_val = entry.get("date") if isinstance(date_val, datetime): d = _normalize_date(date_val) elif isinstance(date_val, str): try: d = _normalize_date(datetime.fromisoformat(date_val)) except Exception: continue else: continue seen_dates.add(d) normalized.append({"date": d, "count": entry.get("count", 0)}) if today in seen_dates: # already recorded today; do not change counts normalized.sort(key=lambda x: x["date"]) if len(normalized) > 32: normalized = normalized[-32:] return normalized # Fill missing days from last recorded date to today normalized.sort(key=lambda x: x["date"]) last_date = normalized[-1]["date"] cursor = last_date + timedelta(days=1) while cursor < today: normalized.append({"date": cursor, "count": 0}) cursor += timedelta(days=1) # Add today with default count=1 normalized.append({"date": today, "count": 1}) # Enforce 32-entry limit (keep most recent 32 by date) normalized.sort(key=lambda x: x["date"]) if len(normalized) > 32: normalized = normalized[-32:] return normalized def get_selfi_short_client(): """Get MongoDB client for the Selfie Short database (DEFAULT)""" if "selfi_short" in mongodb_clients: try: mongodb_clients["selfi_short"].admin.command('ping') return mongodb_clients["selfi_short"] except: del mongodb_clients["selfi_short"] if not config.MONGODB_SELFI_SHORT: print("[MongoDB] MONGODB_SELFI_SHORT connection string not set") return None try: client = MongoClient( config.MONGODB_SELFI_SHORT, tlsAllowInvalidCertificates=True, serverSelectionTimeoutMS=15000, connectTimeoutMS=20000, socketTimeoutMS=30000 ) client.admin.command('ping') mongodb_clients["selfi_short"] = client print(f"[MongoDB] Successfully connected to Selfie Short database") return client except Exception as e: print(f"[MongoDB] Error connecting to Selfie Short MongoDB: {str(e)[:200]}") return None def get_collage_maker_client(): """Get MongoDB client for the collage-maker database""" if "collage_maker" in mongodb_clients: try: mongodb_clients["collage_maker"].admin.command('ping') return mongodb_clients["collage_maker"] except: del mongodb_clients["collage_maker"] if not config.MONGODB_COLLAGE_MAKER: print("[MongoDB] MONGODB_COLLAGE_MAKER connection string not set") return None try: client = MongoClient( config.MONGODB_COLLAGE_MAKER, tlsAllowInvalidCertificates=True, serverSelectionTimeoutMS=15000, connectTimeoutMS=20000, socketTimeoutMS=30000 ) client.admin.command('ping') mongodb_clients["collage_maker"] = client print(f"[MongoDB] Successfully connected to collage-maker database") return client except Exception as e: print(f"[MongoDB] Error connecting to collage-maker MongoDB: {str(e)[:200]}") return None def get_ai_enhancer_client(): """Get MongoDB client for the AI-Enhancer database""" if "ai_enhancer" in mongodb_clients: try: mongodb_clients["ai_enhancer"].admin.command('ping') return mongodb_clients["ai_enhancer"] except: del mongodb_clients["ai_enhancer"] if not config.MONGODB_AI_ENHANCER: print("[MongoDB] MONGODB_AI_ENHANCER connection string not set") return None try: client = MongoClient( config.MONGODB_AI_ENHANCER, tlsAllowInvalidCertificates=True, serverSelectionTimeoutMS=15000, connectTimeoutMS=20000, socketTimeoutMS=30000 ) client.admin.command('ping') mongodb_clients["ai_enhancer"] = client print(f"[MongoDB] Successfully connected to AI-Enhancer database") return client except Exception as e: print(f"[MongoDB] Error connecting to AI-Enhancer MongoDB: {str(e)[:200]}") return None def save_media_click(user_id: Optional[Union[int, str]], category_id: str, appname: Optional[str] = None): """ Save or update media click in the appropriate MongoDB database based on appname. Args: user_id: Optional integer or ObjectId string. If None, a new ObjectId will be generated automatically. category_id: Required MongoDB ObjectId string for category. appname: Optional app name (case-insensitive). - If "collage-maker", uses collage-maker MongoDB (adminPanel database) with "media_clicks" collection - If "AI-Enhancer", uses AI-Enhancer MongoDB (test database) with "media_clicks" collection - Otherwise (DEFAULT), uses Selfie Short MongoDB (adminPanel database) with "media_clicks" collection Document Structure: { userId: ObjectId, categories: [ { categoryId: ObjectId, click_count: int, lastClickedAt: Date } ], ai_edit_complete: int, # Count of times user used any model (default: 0 for new users) ai_edit_last_date: Date, # Last date user used any model (stored as Date object) ai_edit_daily_count: [ # Daily usage tracking (last 32 days) {"date": Date, "count": int}, ... ], createdAt: Date, updatedAt: Date } """ # Log the received appname for debugging print(f"[MongoDB] save_media_click called - appname: '{appname}', user_id: {user_id}, category_id: {category_id}") # Determine which MongoDB to use based on appname (case-insensitive) use_collage_maker = (appname and appname.lower() == "collage-maker") use_ai_enhancer = (appname and appname.lower() == "ai-enhancer") print(f"[MongoDB] Database selection - use_collage_maker: {use_collage_maker}, use_ai_enhancer: {use_ai_enhancer}, default: Selfie Short") # Validate connection string is configured if use_collage_maker: if not config.MONGODB_COLLAGE_MAKER: print("[MongoDB] MONGODB_COLLAGE_MAKER not configured, skipping media_clicks logging") return False elif use_ai_enhancer: if not config.MONGODB_AI_ENHANCER: print("[MongoDB] MONGODB_AI_ENHANCER not configured, skipping media_clicks logging") return False else: # Default: Selfie Short if not config.MONGODB_SELFI_SHORT: print("[MongoDB] MONGODB_SELFI_SHORT not configured, skipping media_clicks logging") return False # Validate category_id if not category_id: print("[MongoDB] category_id not provided, skipping media_clicks logging") return False try: # Get MongoDB client and select database/collection based on appname if use_collage_maker: mongo_client = get_collage_maker_client() if mongo_client is None: print(f"[MongoDB] ERROR: Could not connect to collage-maker MongoDB") if "collage_maker" in mongodb_clients: del mongodb_clients["collage_maker"] return False db = mongo_client.get_database(config.MONGODB_COLLAGE_MAKER_ADMIN_DB_NAME) collection_name = "media_clicks" elif use_ai_enhancer: mongo_client = get_ai_enhancer_client() if mongo_client is None: print(f"[MongoDB] ERROR: Could not connect to AI-Enhancer MongoDB") if "ai_enhancer" in mongodb_clients: del mongodb_clients["ai_enhancer"] return False db = mongo_client.get_database(config.MONGODB_AI_ENHANCER_ADMIN_DB_NAME) collection_name = "media_clicks" print(f"[MongoDB] AI-Enhancer: Using database '{config.MONGODB_AI_ENHANCER_ADMIN_DB_NAME}', collection '{collection_name}'") else: # Default: Selfie Short (selfie-beauty-camera -> adminPanel -> media_clicks) mongo_client = get_selfi_short_client() if mongo_client is None: print(f"[MongoDB] ERROR: Could not connect to Selfie Short MongoDB") if "selfi_short" in mongodb_clients: del mongodb_clients["selfi_short"] return False db = mongo_client.get_database(config.MONGODB_SELFI_SHORT_ADMIN_DB_NAME) collection_name = "media_clicks" print(f"[MongoDB] Selfie Short (DEFAULT): Using database '{config.MONGODB_SELFI_SHORT_ADMIN_DB_NAME}', collection '{collection_name}'") collection = db.get_collection(collection_name) # Handle user_id: use provided ObjectId string directly or convert integer; otherwise generate if user_id is None: user_object_id = ObjectId() print(f"[MongoDB] user_id not provided, generated new ObjectId: {user_object_id}") else: try: user_id_str = str(user_id).strip() if len(user_id_str) == 24: user_object_id = ObjectId(user_id_str) print(f"[MongoDB] Using provided user_id as ObjectId: {user_object_id}") else: # Try integer -> deterministic ObjectId based on hex padded to 24 chars user_id_int = int(user_id_str) user_id_hex = hex(user_id_int)[2:].ljust(24, "0")[:24] user_object_id = ObjectId(user_id_hex) print(f"[MongoDB] Converted integer user_id '{user_id_str}' -> ObjectId: {user_object_id}") except Exception as e: print(f"[MongoDB] Error converting user_id to ObjectId: {e}") user_object_id = ObjectId() print(f"[MongoDB] Generated new ObjectId due to conversion error: {user_object_id}") # Convert category_id to ObjectId try: category_object_id = ObjectId(category_id) except Exception as e: print(f"[MongoDB] Error converting category_id to ObjectId: {e}, category_id: {category_id}") return False now = datetime.utcnow() # Check if document with userId exists existing_doc = collection.find_one({"userId": user_object_id}) if existing_doc: # Build daily counts today_daily_counts = _build_ai_edit_daily_count(existing_doc.get("ai_edit_daily_count"), now) print(f"[MongoDB] Found existing document for userId: {user_object_id}") # Check if category exists in categories array category_exists = False for cat in existing_doc.get("categories", []): cat_id = cat.get("categoryId") if isinstance(cat_id, ObjectId): if cat_id == category_object_id: category_exists = True break elif str(cat_id) == str(category_object_id): category_exists = True break if category_exists: # Category exists, increment click_count result = collection.update_one( { "userId": user_object_id, "categories.categoryId": category_object_id }, { "$inc": { "categories.$.click_count": 1, "ai_edit_complete": 1 }, "$set": { "categories.$.lastClickedAt": now, "ai_edit_last_date": now, "updatedAt": now, "ai_edit_daily_count": today_daily_counts } } ) print(f"[MongoDB] Updated category click_count - userId: {user_object_id}, categoryId: {category_object_id}") else: # Category doesn't exist, add new category to array result = collection.update_one( {"userId": user_object_id}, { "$push": { "categories": { "categoryId": category_object_id, "click_count": 1, "lastClickedAt": now } }, "$inc": { "ai_edit_complete": 1 }, "$set": { "ai_edit_last_date": now, "updatedAt": now, "ai_edit_daily_count": today_daily_counts } } ) print(f"[MongoDB] Added new category - userId: {user_object_id}, categoryId: {category_object_id}") else: # Document doesn't exist, create new document today_daily_counts_new = _build_ai_edit_daily_count(None, now) result = collection.find_one_and_update( {"userId": user_object_id}, { "$setOnInsert": { "userId": user_object_id, "createdAt": now, "ai_edit_complete": 1, "ai_edit_daily_count": today_daily_counts_new }, "$push": { "categories": { "categoryId": category_object_id, "click_count": 1, "lastClickedAt": now } }, "$set": { "ai_edit_last_date": now, "updatedAt": now } }, upsert=True, return_document=ReturnDocument.AFTER ) print(f"[MongoDB] Created new document - userId: {user_object_id}, categoryId: {category_object_id}") return True except Exception as e: import traceback print(f"[MongoDB] ERROR saving media_clicks: {str(e)}") print(f"[MongoDB] Traceback: {traceback.format_exc()}") # Clear cached client on error so next attempt will retry if use_collage_maker: if "collage_maker" in mongodb_clients: del mongodb_clients["collage_maker"] elif use_ai_enhancer: if "ai_enhancer" in mongodb_clients: del mongodb_clients["ai_enhancer"] else: if "selfi_short" in mongodb_clients: del mongodb_clients["selfi_short"] return False def get_category_name(category_id: str, appname: Optional[str] = None) -> Optional[str]: """ Look up category name from the appropriate MongoDB based on appname -> admin DB -> category collection. Args: category_id: MongoDB ObjectId string for the category appname: App name to determine which MongoDB to use - "collage-maker" → MONGODB_COLLAGE_MAKER - "AI-Enhancer" → MONGODB_AI_ENHANCER - default → MONGODB_SELFI_SHORT Returns: Category name string or None if not found """ if not category_id: return None # Determine which MongoDB to use based on appname (case-insensitive) use_collage_maker = (appname and appname.lower() == "collage-maker") use_ai_enhancer = (appname and appname.lower() == "ai-enhancer") try: # Get the appropriate MongoDB client based on appname if use_collage_maker: if not config.MONGODB_COLLAGE_MAKER: print("[MongoDB] MONGODB_COLLAGE_MAKER not configured, cannot lookup category name") return None mongo_client = get_collage_maker_client() admin_db_name = config.MONGODB_COLLAGE_MAKER_ADMIN_DB_NAME # adminPanel elif use_ai_enhancer: if not config.MONGODB_AI_ENHANCER: print("[MongoDB] MONGODB_AI_ENHANCER not configured, cannot lookup category name") return None mongo_client = get_ai_enhancer_client() admin_db_name = config.MONGODB_AI_ENHANCER_ADMIN_DB_NAME # test else: # Default: Selfie Short if not config.MONGODB_SELFI_SHORT: print("[MongoDB] MONGODB_SELFI_SHORT not configured, cannot lookup category name") return None mongo_client = get_selfi_short_client() admin_db_name = config.MONGODB_SELFI_SHORT_ADMIN_DB_NAME # adminPanel if mongo_client is None: print(f"[MongoDB] Could not connect to MongoDB for category lookup") return None # Access adminPanel database -> categories collection db = mongo_client.get_database(admin_db_name) collection = db.get_collection("categories") # Convert category_id to ObjectId try: category_object_id = ObjectId(category_id) except Exception as e: print(f"[MongoDB] Error converting category_id to ObjectId: {e}") return None # Find the category document category_doc = collection.find_one({"_id": category_object_id}) if category_doc: category_name = category_doc.get("name") print(f"[MongoDB] Found category name: {category_name} for id: {category_id} (appname: {appname})") return category_name else: print(f"[MongoDB] Category not found for id: {category_id}") return None except Exception as e: print(f"[MongoDB] Error looking up category name: {str(e)}") return None # Pixverse category mapping - keyword to category name # 4 types of videos: AI BF, AI GF, AI Hug Video, AI Kiss Video PIXVERSE_CATEGORIES = { "kiss": "Ai Kiss Video", "hug": "Ai Hug Video", "boyfriend": "Ai BF", "bf": "Ai BF", "girlfriend": "Ai GF", "gf": "Ai GF" } def get_category_by_prompt(prompt_text: str, appname: Optional[str] = None) -> tuple: """ Match prompt_text against Pixverse category names and get category_id. Args: prompt_text: The prompt text to match against categories appname: App name to determine which MongoDB to use Returns: Tuple of (category_id, category_name) or (None, None) if not found """ if not prompt_text: return None, None prompt_lower = prompt_text.lower() matched_category_name = None # Find matching category based on keywords in prompt for keyword, category_name in PIXVERSE_CATEGORIES.items(): if keyword in prompt_lower: matched_category_name = category_name print(f"[MongoDB] Matched prompt keyword '{keyword}' to category: {category_name}") break if not matched_category_name: print(f"[MongoDB] No category match found for prompt: {prompt_text[:50]}...") return None, None # Determine which MongoDB to use based on appname (case-insensitive) use_collage_maker = (appname and appname.lower() == "collage-maker") use_ai_enhancer = (appname and appname.lower() == "ai-enhancer") try: # Get the appropriate MongoDB client based on appname if use_collage_maker: if not config.MONGODB_COLLAGE_MAKER: print("[MongoDB] MONGODB_COLLAGE_MAKER not configured") return None, matched_category_name mongo_client = get_collage_maker_client() admin_db_name = config.MONGODB_COLLAGE_MAKER_ADMIN_DB_NAME elif use_ai_enhancer: if not config.MONGODB_AI_ENHANCER: print("[MongoDB] MONGODB_AI_ENHANCER not configured") return None, matched_category_name mongo_client = get_ai_enhancer_client() admin_db_name = config.MONGODB_AI_ENHANCER_ADMIN_DB_NAME else: # Default: Selfie Short if not config.MONGODB_SELFI_SHORT: print("[MongoDB] MONGODB_SELFI_SHORT not configured") return None, matched_category_name mongo_client = get_selfi_short_client() admin_db_name = config.MONGODB_SELFI_SHORT_ADMIN_DB_NAME if mongo_client is None: print(f"[MongoDB] Could not connect to MongoDB for category lookup by prompt") return None, matched_category_name # Access adminPanel database -> categories collection db = mongo_client.get_database(admin_db_name) collection = db.get_collection("categories") # Find the category document by name category_doc = collection.find_one({"name": matched_category_name}) if category_doc: category_id = str(category_doc.get("_id")) print(f"[MongoDB] Found category_id: {category_id} for name: {matched_category_name}") return category_id, matched_category_name else: print(f"[MongoDB] Category not found in database for name: {matched_category_name}") return None, matched_category_name except Exception as e: print(f"[MongoDB] Error looking up category by prompt: {str(e)}") return None, matched_category_name def get_logs_client(): """Get MongoDB client for the Logs database""" if "logs" in mongodb_clients: try: mongodb_clients["logs"].admin.command('ping') return mongodb_clients["logs"] except: del mongodb_clients["logs"] if not config.MONGODB_LOGS: print("[MongoDB] MONGODB_LOGS connection string not set") return None try: client = MongoClient( config.MONGODB_LOGS, tlsAllowInvalidCertificates=True, serverSelectionTimeoutMS=15000, connectTimeoutMS=20000, socketTimeoutMS=30000 ) client.admin.command('ping') mongodb_clients["logs"] = client print(f"[MongoDB] Successfully connected to Logs database") return client except Exception as e: print(f"[MongoDB] Error connecting to Logs MongoDB: {str(e)[:200]}") return None def save_request_log( user_id: Optional[str], subcategory: str, endpoint: str, status: str, response_time: float, model: str = "PixVerse", appname: Optional[str] = None, error: Optional[str] = None ): """ Save request log to MongoDB Logs database. Args: user_id: User ID string subcategory: Category/subcategory name endpoint: API endpoint called status: "success" or "error" response_time: Time taken in seconds model: Model name (default: "PixVerse") appname: App name if provided error: Error message if status is "error" Document Structure: { user_id: string, subcategory: string, endpoint: string, status: string, response_time: float, model: string, timestamp: Date, appname: string or null, error: string or null } """ if not config.MONGODB_LOGS: print("[MongoDB] MONGODB_LOGS not configured, skipping request logging") return False try: mongo_client = get_logs_client() if mongo_client is None: print(f"[MongoDB] ERROR: Could not connect to Logs MongoDB") if "logs" in mongodb_clients: del mongodb_clients["logs"] return False db = mongo_client.get_database(config.MONGODB_LOGS_DB_NAME) collection = db.get_collection(config.MONGODB_LOGS_COLLECTION) log_doc = { "user_id": user_id, "subcategory": subcategory, "endpoint": endpoint, "status": status, "response_time": response_time, "model": model, "timestamp": datetime.utcnow(), "appname": appname, "error": error } result = collection.insert_one(log_doc) print(f"[MongoDB] Request log saved - user_id: {user_id}, status: {status}, endpoint: {endpoint}") return True except Exception as e: import traceback print(f"[MongoDB] ERROR saving request log: {str(e)}") print(f"[MongoDB] Traceback: {traceback.format_exc()}") if "logs" in mongodb_clients: del mongodb_clients["logs"] return False