storage/mongo_storage.py

"""
storage/mongo_storage.py — MongoDB Atlas: quản lý 4 collections theo ERD

Collections:
    raw_images           — Ảnh gốc từ camera / upload
    face_crops           — Ảnh khuôn mặt đã crop (có đầy đủ metadata)
    unknown_faces        — Khuôn mặt không nhận dạng được
    recognition_payloads — Payload debug cho mỗi lần nhận diện
"""
import logging
import time
from datetime import datetime, timezone

from bson.objectid import ObjectId
from pymongo import MongoClient

import config

logger = logging.getLogger("mongo_storage")


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


class MongoStorage:
    """Singleton wrapper cho tất cả thao tác MongoDB."""

    def __init__(self):
        self._client: MongoClient | None = None

    def _get_db(self):
        if self._client is None:
            self._client = MongoClient(config.MONGO_URI)
        return self._client[config.MONGO_DB_NAME]

    # ── raw_images ────────────────────────────────────────────────────────

    def save_raw_image(
        self,
        img_bytes: bytes,
        source_type: str = "upload",
        width: int = 0,
        height: int = 0,
        fmt: str = "jpeg",
    ) -> str:
        """
        Lưu ảnh gốc vào collection raw_images.
        Returns:
            str: ObjectId string của document vừa lưu (rỗng nếu lỗi)
        """
        try:
            col = self._get_db()[config.MONGO_COL_RAW_IMAGES]
            doc = {
                "image_bytes": img_bytes,
                "source_type": source_type,  # "webcam" | "upload"
                "camera_id": None,
                "width": width,
                "height": height,
                "format": fmt,
                "metadata": {},
                "created_at": _utcnow(),
            }
            result = col.insert_one(doc)
            logger.debug(f"[Mongo] Saved raw_image: {result.inserted_id}")
            return str(result.inserted_id)
        except Exception as e:
            logger.error(f"[Mongo] save_raw_image failed: {e}")
            return ""

    def get_raw_image(self, raw_image_id: str) -> bytes | None:
        """Lấy bytes ảnh gốc theo ID."""
        try:
            col = self._get_db()[config.MONGO_COL_RAW_IMAGES]
            doc = col.find_one({"_id": ObjectId(raw_image_id)})
            return doc.get("image_bytes") if doc else None
        except Exception as e:
            logger.error(f"[Mongo] get_raw_image failed: {e}")
            return None

    # ── face_crops ────────────────────────────────────────────────────────

    def save_face_crop(
        self,
        img_bytes: bytes,
        raw_image_id: str = "",
        person_id: str | None = None,
        visit_id: str | None = None,
        recognition_log_id: str | None = None,
        bbox: dict | None = None,
        image_type: str = "recognition",
        confidence: float = 0.0,
        quality_score: float = 0.0,
        model_name: str = "buffalo_l",
        model_version: str = "1.0.0",
    ) -> str:
        """
        Lưu ảnh khuôn mặt crop vào collection face_crops (theo ERD).
        Returns:
            str: ObjectId string (rỗng nếu lỗi)
        """
        try:
            col = self._get_db()[config.MONGO_COL_FACE_CROPS]
            doc = {
                # Foreign keys
                "person_id": person_id,
                "visit_id": visit_id,
                "recognition_log_id": recognition_log_id,
                "raw_image_id": raw_image_id,
                # Image data
                "image_bytes": img_bytes,
                "image_url": None,
                "image_base64": None,
                # Bounding box & detection info
                "bbox": bbox or {},
                "image_type": image_type,  # "recognition" | "registration"
                "confidence": confidence,
                "quality_score": quality_score,
                # Model info
                "model": {
                    "name": model_name,
                    "version": model_version,
                },
                "created_at": _utcnow(),
            }
            result = col.insert_one(doc)
            logger.debug(f"[Mongo] Saved face_crop: {result.inserted_id}")
            return str(result.inserted_id)
        except Exception as e:
            logger.error(f"[Mongo] save_face_crop failed: {e}")
            return ""

    def update_face_crop_links(
        self,
        crop_id: str,
        person_id: str,
        visit_id: str,
        recognition_log_id: str,
    ) -> bool:
        """Cập nhật các liên kết sau khi nhận diện thành công."""
        try:
            col = self._get_db()[config.MONGO_COL_FACE_CROPS]
            result = col.update_one(
                {"_id": ObjectId(crop_id)},
                {"$set": {
                    "person_id": person_id,
                    "visit_id": visit_id,
                    "recognition_log_id": recognition_log_id,
                }},
            )
            return result.modified_count > 0
        except Exception as e:
            logger.error(f"[Mongo] update_face_crop_links failed: {e}")
            return False

    def get_face_crop_bytes(self, crop_id: str) -> bytes | None:
        """Lấy bytes ảnh crop theo ID."""
        try:
            col = self._get_db()[config.MONGO_COL_FACE_CROPS]
            doc = col.find_one({"_id": ObjectId(crop_id)})
            return doc.get("image_bytes") if doc else None
        except Exception as e:
            logger.error(f"[Mongo] get_face_crop_bytes failed: {e}")
            return None

    def delete_face_crop(self, crop_id: str) -> bool:
        """Xóa ảnh crop theo ID."""
        try:
            col = self._get_db()[config.MONGO_COL_FACE_CROPS]
            result = col.delete_one({"_id": ObjectId(crop_id)})
            return result.deleted_count > 0
        except Exception as e:
            logger.error(f"[Mongo] delete_face_crop failed: {e}")
            return False

    def list_face_crops_by_person(self, person_id: str) -> list[str]:
        """Lấy danh sách crop_id của một person."""
        try:
            col = self._get_db()[config.MONGO_COL_FACE_CROPS]
            docs = col.find({"person_id": person_id}, {"_id": 1})
            return [str(d["_id"]) for d in docs]
        except Exception as e:
            logger.error(f"[Mongo] list_face_crops_by_person failed: {e}")
            return []

    # ── unknown_faces ─────────────────────────────────────────────────────

    def save_unknown_face(
        self,
        face_crop_id: str,
        raw_image_id: str,
        recognition_log_id: str,
        reason: str,
        confidence: float,
        similarity_distance: float,
    ) -> str:
        """
        Lưu thông tin khuôn mặt không nhận diện được vào unknown_faces.
        Returns:
            str: ObjectId string (rỗng nếu lỗi)
        """
        try:
            col = self._get_db()[config.MONGO_COL_UNKNOWN_FACES]
            doc = {
                "face_crop_id": face_crop_id,
                "raw_image_id": raw_image_id,
                "recognition_log_id": recognition_log_id,
                "reason": reason,            # "no_match" | "low_confidence" | "no_db"
                "confidence": confidence,
                "similarity_distance": similarity_distance,
                "is_reviewed": False,
                "review_result": None,
                "reviewed_by": None,
                "reviewed_at": None,
                "review_note": None,
                "created_at": _utcnow(),
            }
            result = col.insert_one(doc)
            logger.debug(f"[Mongo] Saved unknown_face: {result.inserted_id}")
            return str(result.inserted_id)
        except Exception as e:
            logger.error(f"[Mongo] save_unknown_face failed: {e}")
            return ""

    def list_unknown_faces(self, limit: int = 50) -> list[dict]:
        """Lấy danh sách unknown faces gần nhất."""
        try:
            col = self._get_db()[config.MONGO_COL_UNKNOWN_FACES]
            docs = col.find({}, {"image_bytes": 0}).sort("created_at", -1).limit(limit)
            results = []
            for d in docs:
                d["id"] = str(d.pop("_id"))
                if "created_at" in d:
                    d["created_at"] = d["created_at"].isoformat()
                results.append(d)
            return results
        except Exception as e:
            logger.error(f"[Mongo] list_unknown_faces failed: {e}")
            return []

    # ── recognition_payloads ──────────────────────────────────────────────

    def save_recognition_payload(
        self,
        recognition_log_id: str,
        raw_image_id: str,
        face_crop_id: str,
        model_name: str,
        model_version: str,
        request_payload: dict,
        response_payload: dict,
        runtime_ms: float,
    ) -> str:
        """
        Lưu payload nhận diện đầy đủ vào recognition_payloads (để debug).
        Returns:
            str: ObjectId string (rỗng nếu lỗi)
        """
        try:
            col = self._get_db()[config.MONGO_COL_PAYLOADS]
            doc = {
                "recognition_log_id": recognition_log_id,
                "raw_image_id": raw_image_id,
                "face_crop_id": face_crop_id,
                "model_name": model_name,
                "model_version": model_version,
                "request_payload": request_payload,
                "response_payload": response_payload,
                "runtime": {"total_ms": runtime_ms},
                "created_at": _utcnow(),
            }
            result = col.insert_one(doc)
            logger.debug(f"[Mongo] Saved recognition_payload: {result.inserted_id}")
            return str(result.inserted_id)
        except Exception as e:
            logger.error(f"[Mongo] save_recognition_payload failed: {e}")
            return ""


# ── Singleton ─────────────────────────────────────────────────────────────
mongo_storage = MongoStorage()