Add logging status endpoint and improved admin Mongo initialization logging

api/main.py

@@ -2,17 +2,28 @@
 import os
 import uuid
 import shutil
+import re
 from datetime import datetime
 from typing import Dict, List, Optional
 
 import numpy as np
-from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, Header, Request
+from fastapi import (
+    FastAPI,
+    UploadFile,
+    File,
+    HTTPException,
+    Depends,
+    Header,
+    Request,
+    Form,
+)
 from fastapi.responses import FileResponse, JSONResponse
 from pydantic import BaseModel
 from PIL import Image
 import cv2
 import logging
 
+from bson import ObjectId
 from pymongo import MongoClient
 import time
 
@@ -47,6 +58,54 @@ mongo_client = MongoClient(MONGO_URI)
 mongo_db = mongo_client["object_remover"]
 mongo_logs = mongo_db["api_logs"]
 
+ADMIN_MONGO_URI = os.environ.get("MONGODB_ADMIN")
+DEFAULT_CATEGORY_ID = "69368f722e46bd68ae188984"
+admin_media_clicks = None
+
+
+def _init_admin_mongo() -> None:
+    global admin_media_clicks
+    if not ADMIN_MONGO_URI:
+        log.info("Admin Mongo URI not provided; media click logging disabled")
+        return
+    try:
+        admin_client = MongoClient(ADMIN_MONGO_URI)
+        try:
+            admin_db = admin_client.get_default_database() or admin_client["admin"]
+        except Exception:
+            admin_db = admin_client["admin"]
+        admin_media_clicks = admin_db["media_clicks"]
+        try:
+            admin_media_clicks.drop_index("user_id_1_header_1_media_id_1")
+            log.info("Dropped legacy index user_id_1_header_1_media_id_1")
+        except Exception as idx_err:
+            log.info("Skipping legacy index drop: %s", idx_err)
+        log.info(
+            "Admin media click logging enabled on db=%s collection=%s",
+            admin_media_clicks.database.name,
+            admin_media_clicks.name,
+        )
+    except Exception as err:
+        log.error("Failed to init admin Mongo client: %s", err)
+        admin_media_clicks = None
+
+
+_init_admin_mongo()
+
+
+def _admin_logging_status() -> Dict[str, object]:
+    if admin_media_clicks is None:
+        return {
+            "enabled": False,
+            "db": None,
+            "collection": None,
+        }
+    return {
+        "enabled": True,
+        "db": admin_media_clicks.database.name,
+        "collection": admin_media_clicks.name,
+    }
+
 def bearer_auth(authorization: Optional[str] = Header(default=None)) -> None:
     if not ENV_TOKEN:
         return
@@ -62,12 +121,92 @@ class InpaintRequest(BaseModel):
     mask_id: str
     invert_mask: bool = True  # True => selected/painted area is removed
     passthrough: bool = False  # If True, return the original image unchanged
+    user_id: Optional[str] = None
+    category_id: Optional[str] = None
 
 
 class SimpleRemoveRequest(BaseModel):
     image_id: str  # Image with pink/magenta segments to remove
 
 
+def _coerce_object_id(value: Optional[str]) -> ObjectId:
+    if value is None:
+        return ObjectId()
+    value_str = str(value).strip()
+    if re.fullmatch(r"[0-9a-fA-F]{24}", value_str):
+        return ObjectId(value_str)
+    if value_str.isdigit():
+        hex_str = format(int(value_str), "x")
+        if len(hex_str) > 24:
+            hex_str = hex_str[-24:]
+        hex_str = hex_str.rjust(24, "0")
+        return ObjectId(hex_str)
+    return ObjectId()
+
+
+def _coerce_category_id(category_id: Optional[str]) -> ObjectId:
+    raw = category_id or DEFAULT_CATEGORY_ID
+    raw_str = str(raw).strip()
+    if re.fullmatch(r"[0-9a-fA-F]{24}", raw_str):
+        return ObjectId(raw_str)
+    return _coerce_object_id(raw_str)
+
+
+def log_media_click(user_id: Optional[str], category_id: Optional[str]) -> None:
+    if admin_media_clicks is None:
+        return
+    try:
+        user_obj = _coerce_object_id(user_id)
+        category_obj = _coerce_category_id(category_id)
+        now = datetime.utcnow()
+
+        doc = admin_media_clicks.find_one({"userId": user_obj})
+        if doc:
+            categories = doc.get("categories") or []
+            if any(cat.get("categoryId") == category_obj for cat in categories):
+                admin_media_clicks.update_one(
+                    {"_id": doc["_id"], "categories.categoryId": category_obj},
+                    {
+                        "$inc": {"categories.$.click_count": 1},
+                        "$set": {
+                            "categories.$.lastClickedAt": now,
+                            "updatedAt": now,
+                        },
+                    },
+                )
+            else:
+                admin_media_clicks.update_one(
+                    {"_id": doc["_id"]},
+                    {
+                        "$push": {
+                            "categories": {
+                                "categoryId": category_obj,
+                                "click_count": 1,
+                                "lastClickedAt": now,
+                            }
+                        },
+                        "$set": {"updatedAt": now},
+                    },
+                )
+        else:
+            admin_media_clicks.insert_one(
+                {
+                    "userId": user_obj,
+                    "categories": [
+                        {
+                            "categoryId": category_obj,
+                            "click_count": 1,
+                            "lastClickedAt": now,
+                        }
+                    ],
+                    "createdAt": now,
+                    "updatedAt": now,
+                }
+            )
+    except Exception as err:
+        log.warning("Admin media click logging failed: %s", err)
+
+
 @app.get("/")
 def root() -> Dict[str, object]:
     return {
@@ -93,6 +232,12 @@ def health() -> Dict[str, str]:
     return {"status": "healthy"}
 
 
+@app.get("/logging-status")
+def logging_status(_: None = Depends(bearer_auth)) -> Dict[str, object]:
+    """Helper endpoint to verify admin media logging wiring (no secrets exposed)."""
+    return _admin_logging_status()
+
+
 @app.post("/upload-image")
 def upload_image(image: UploadFile = File(...), _: None = Depends(bearer_auth)) -> Dict[str, str]:
     ext = os.path.splitext(image.filename)[1] or ".png"
@@ -223,6 +368,7 @@ def inpaint(req: InpaintRequest, _: None = Depends(bearer_auth)) -> Dict[str, st
             output_path, "PNG", optimize=False, compress_level=1
         )
 
+        log_media_click(req.user_id, req.category_id)
         return {"result": output_name}
 
     except Exception as e:
@@ -301,6 +447,7 @@ def inpaint_url(req: InpaintRequest, request: Request, _: None = Depends(bearer_
 
     url = str(request.url_for("download_file", filename=result_name))
     logs.append({"result": result_name, "url": url, "timestamp": datetime.utcnow().isoformat()})
+    log_media_click(req.user_id, req.category_id)
     return {"result": result_name, "url": url}
 
 
@@ -312,6 +459,8 @@ def inpaint_multipart(
     invert_mask: bool = True,
     mask_is_painted: bool = False,  # if True, mask file is the painted-on image (e.g., black strokes on original)
     passthrough: bool = False,
+    user_id: Optional[str] = Form(None),
+    category_id: Optional[str] = Form(None),
     _: None = Depends(bearer_auth),
 ) -> Dict[str, str]:
     # Load in-memory
@@ -339,6 +488,7 @@ def inpaint_multipart(
         resp: Dict[str, str] = {"result": result_name}
         if url:
             resp["url"] = url
+        log_media_click(user_id, category_id)
         return resp
 
     if mask_is_painted:
@@ -434,6 +584,7 @@ def inpaint_multipart(
     resp: Dict[str, str] = {"result": result_name}
     if url:
         resp["url"] = url
+    log_media_click(user_id, category_id)
     return resp
 
 
@@ -441,6 +592,8 @@ def inpaint_multipart(
 def remove_pink_segments(
     image: UploadFile = File(...),
     request: Request = None,
+    user_id: Optional[str] = Form(None),
+    category_id: Optional[str] = Form(None),
     _: None = Depends(bearer_auth),
 ) -> Dict[str, str]:
     """
@@ -545,6 +698,7 @@ def remove_pink_segments(
     resp: Dict[str, str] = {"result": result_name, "pink_segments_detected": str(nonzero)}
     if url:
         resp["url"] = url
+    log_media_click(user_id, category_id)
     return resp