storage/pg_storage.py

"""
storage/pg_storage.py — Supabase PostgreSQL: quản lý 4 bảng theo ERD

Tables:
    persons           — Thông tin cá nhân
    visits            — Lượt ghé thăm
    face_embeddings   — Vector embedding 512 chiều
    recognition_logs  — Nhật ký nhận diện (đầy đủ field theo ERD)
"""
import json
import logging
import uuid

import psycopg2
from psycopg2.extras import RealDictCursor

import config

logger = logging.getLogger("pg_storage")


class PostgresStorage:
    """Quản lý tất cả thao tác SQL trên Supabase PostgreSQL."""

    def __init__(self, uri: str = config.SUPABASE_DB_URI):
        self.uri = uri

    def _conn(self):
        return psycopg2.connect(self.uri)

    # ── Load embeddings ──────────────────────────────────────────────────

    def load_all_embeddings(self) -> list[dict]:
        """
        Nạp toàn bộ active embeddings và metadata người dùng từ DB.
        Trả về list[{person_id, full_name, embedding, face_crop_mongo_id, embedding_id}]
        """
        query = """
            SELECT
                p.id::text        AS person_id,
                p.full_name,
                e.id::text        AS embedding_id,
                e.embedding,
                e.face_crop_mongo_id
            FROM persons p
            JOIN face_embeddings e ON p.id = e.person_id
            WHERE e.is_active = TRUE
            ORDER BY p.full_name;
        """
        conn = self._conn()
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute(query)
                rows = cur.fetchall()
                result = []
                for r in rows:
                    emb = r["embedding"]
                    if isinstance(emb, str):
                        cleaned = emb.strip("[]")
                        embedding = [float(x) for x in cleaned.split(",") if x.strip()]
                    elif isinstance(emb, list):
                        embedding = [float(x) for x in emb]
                    else:
                        continue
                    result.append({
                        "person_id":        r["person_id"],
                        "full_name":        r["full_name"],
                        "embedding_id":     r["embedding_id"],
                        "embedding":        embedding,
                        "face_crop_mongo_id": r["face_crop_mongo_id"],
                    })
                logger.info(f"[PG] Loaded {len(result)} embeddings from Supabase.")
                return result
        finally:
            conn.close()

    # ── match_face RPC ───────────────────────────────────────────────────

    def match_face_rpc(
        self,
        embedding: list[float],
        threshold: float = config.SIMILARITY_THRESHOLD,
    ) -> dict | None:
        """
        Gọi Supabase RPC function `match_face` để tìm khuôn mặt khớp nhất.
        Returns:
            dict với {person_id, full_name, embedding_id, similarity} hoặc None
        """
        emb_str = "[" + ",".join(map(str, embedding)) + "]"
        conn = self._conn()
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM match_face(%s::vector, %s);",
                    (emb_str, threshold),
                )
                row = cur.fetchone()
                if row:
                    return dict(row)
                return None
        except Exception as e:
            logger.warning(f"[PG] match_face RPC failed (will fallback to RAM): {e}")
            return None
        finally:
            conn.close()

    # ── register_new_face ────────────────────────────────────────────────

    def register_new_face(
        self,
        name: str,
        embedding: list[float],
        face_crop_mongo_id: str,
        model_name: str = config.MODEL_NAME,
    ) -> tuple[str, str]:
        """
        Đăng ký người mới: thêm vào persons + face_embeddings.
        Returns:
            (person_id, embedding_id)
        """
        person_id = str(uuid.uuid4())
        emb_id    = str(uuid.uuid4())
        emb_str   = "[" + ",".join(map(str, embedding)) + "]"

        conn = self._conn()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO persons
                        (id, full_name, status, total_appearances, total_visits, created_at, updated_at)
                    VALUES (%s, %s, 'active', 1, 1, NOW(), NOW());
                    """,
                    (person_id, name),
                )
                cur.execute(
                    """
                    INSERT INTO face_embeddings
                        (id, person_id, embedding, model_name, model_version,
                         face_crop_mongo_id, is_primary, is_active, created_at)
                    VALUES (%s, %s, %s, %s, '1.0.0', %s, TRUE, TRUE, NOW());
                    """,
                    (emb_id, person_id, emb_str, model_name, face_crop_mongo_id),
                )
                conn.commit()
                logger.info(f"[PG] Registered new face: {name} (person_id={person_id})")
                return person_id, emb_id
        except Exception as e:
            conn.rollback()
            logger.error(f"[PG] register_new_face failed: {e}")
            raise
        finally:
            conn.close()

    # ── record_matched_visit ─────────────────────────────────────────────

    def record_matched_visit(
        self,
        person_id: str,
        embedding_id: str,
        face_crop_mongo_id: str,
        raw_image_mongo_id: str,
        similarity: float,
        bbox: dict,
        ui_color: str = "#00DC50",
        ui_track_index: int = 0,
    ) -> tuple[str, str]:
        """
        Ghi nhận lượt ghé thăm cho khuôn mặt khớp.
        - Tạo hoặc cập nhật visit (trong 30 phút)
        - Ghi recognition_logs đầy đủ field theo ERD
        - Cập nhật persons (total_appearances, total_visits, last_seen_at)
        Returns:
            (visit_id, log_id)
        """
        log_id = str(uuid.uuid4())
        visit_id = None

        conn = self._conn()
        try:
            with conn.cursor() as cur:
                # 1. Kiểm tra visit trong 30 phút gần nhất
                cur.execute(
                    """
                    SELECT id FROM visits
                    WHERE person_id = %s
                      AND last_seen_at >= NOW() - INTERVAL '30 minutes'
                    ORDER BY last_seen_at DESC LIMIT 1;
                    """,
                    (person_id,),
                )
                row = cur.fetchone()
                if row:
                    visit_id = str(row[0])
                    cur.execute(
                        """
                        UPDATE visits
                        SET last_seen_at   = NOW(),
                            appearance_count = appearance_count + 1
                        WHERE id = %s;
                        """,
                        (visit_id,),
                    )
                else:
                    visit_id = str(uuid.uuid4())
                    cur.execute(
                        """
                        INSERT INTO visits
                            (id, person_id, first_seen_at, last_seen_at, appearance_count, created_at)
                        VALUES (%s, %s, NOW(), NOW(), 1, NOW());
                        """,
                        (visit_id, person_id),
                    )

                # 2. Ghi recognition_log đầy đủ field theo ERD
                cur.execute(
                    """
                    INSERT INTO recognition_logs (
                        id, person_id, visit_id, embedding_id,
                        raw_image_mongo_id, face_crop_mongo_id,
                        recognition_status, similarity_distance, confidence,
                        bbox_x, bbox_y, bbox_width, bbox_height,
                        ui_color, ui_track_index,
                        detected_at
                    )
                    VALUES (
                        %s, %s, %s, %s,
                        %s, %s,
                        'recognized', %s, %s,
                        %s, %s, %s, %s,
                        %s, %s,
                        NOW()
                    );
                    """,
                    (
                        log_id, person_id, visit_id, embedding_id,
                        raw_image_mongo_id, face_crop_mongo_id,
                        float(1.0 - similarity), float(similarity),
                        float(bbox.get("x1", 0)),
                        float(bbox.get("y1", 0)),
                        float(bbox.get("x2", 0) - bbox.get("x1", 0)),
                        float(bbox.get("y2", 0) - bbox.get("y1", 0)),
                        ui_color, ui_track_index,
                    ),
                )

                # 3. Cập nhật persons
                cur.execute(
                    """
                    UPDATE persons
                    SET total_appearances = total_appearances + 1,
                        total_visits      = (SELECT COUNT(*) FROM visits WHERE person_id = %s),
                        last_seen_at      = NOW(),
                        updated_at        = NOW()
                    WHERE id = %s;
                    """,
                    (person_id, person_id),
                )

                conn.commit()
                logger.info(f"[PG] Matched visit recorded: person={person_id}, visit={visit_id}")
                return visit_id, log_id

        except Exception as e:
            conn.rollback()
            logger.error(f"[PG] record_matched_visit failed: {e}")
            raise
        finally:
            conn.close()

    # ── record_unknown_log ───────────────────────────────────────────────

    def record_unknown_log(
        self,
        face_crop_mongo_id: str,
        raw_image_mongo_id: str,
        similarity: float,
        bbox: dict,
        ui_color: str = "#FF5500",
        ui_track_index: int = 0,
    ) -> str:
        """
        Ghi recognition_log với person_id = NULL cho khuôn mặt không nhận diện.
        Returns:
            log_id (str)
        """
        log_id = str(uuid.uuid4())
        conn = self._conn()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO recognition_logs (
                        id, person_id, visit_id, embedding_id,
                        raw_image_mongo_id, face_crop_mongo_id,
                        recognition_status, similarity_distance, confidence,
                        bbox_x, bbox_y, bbox_width, bbox_height,
                        ui_color, ui_track_index,
                        detected_at
                    )
                    VALUES (
                        %s, NULL, NULL, NULL,
                        %s, %s,
                        'unknown', %s, %s,
                        %s, %s, %s, %s,
                        %s, %s,
                        NOW()
                    );
                    """,
                    (
                        log_id,
                        raw_image_mongo_id, face_crop_mongo_id,
                        float(1.0 - similarity), float(similarity),
                        float(bbox.get("x1", 0)),
                        float(bbox.get("y1", 0)),
                        float(bbox.get("x2", 0) - bbox.get("x1", 0)),
                        float(bbox.get("y2", 0) - bbox.get("y1", 0)),
                        ui_color, ui_track_index,
                    ),
                )
                conn.commit()
                logger.info(f"[PG] Unknown face logged: log_id={log_id}")
                return log_id
        except Exception as e:
            conn.rollback()
            logger.error(f"[PG] record_unknown_log failed: {e}")
            raise
        finally:
            conn.close()

    # ── update_person_info ───────────────────────────────────────────────

    def update_person_info(
        self,
        person_id: str,
        full_name: str,
        date_of_birth: str | None = None,
        phone_number: str | None = None,
        address: str | None = None,
    ) -> bool:
        """Cập nhật thông tin cá nhân của một person."""
        conn = self._conn()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    UPDATE persons
                    SET full_name    = %s,
                        date_of_birth = %s,
                        phone_number  = %s,
                        address       = %s,
                        status        = 'active',
                        updated_at    = NOW()
                    WHERE id = %s;
                    """,
                    (full_name, date_of_birth, phone_number, address, person_id),
                )
                updated = cur.rowcount > 0
                conn.commit()
                return updated
        except Exception as e:
            conn.rollback()
            logger.error(f"[PG] update_person_info failed: {e}")
            raise
        finally:
            conn.close()

    # ── delete_person ────────────────────────────────────────────────────

    def delete_person(self, person_id: str) -> tuple[bool, list[str]]:
        """
        Xóa person và toàn bộ dữ liệu liên quan trên Supabase.
        Returns:
            (success, list of face_crop_mongo_ids cần xóa trên MongoDB)
        """
        conn = self._conn()
        try:
            with conn.cursor() as cur:
                cur.execute("SELECT id FROM persons WHERE id = %s;", (person_id,))
                if not cur.fetchone():
                    return False, []

                # Thu thập tất cả mongo_ids cần xóa
                mongo_ids: list[str] = []
                cur.execute(
                    "SELECT face_crop_mongo_id FROM face_embeddings WHERE person_id = %s;",
                    (person_id,),
                )
                mongo_ids.extend([r[0] for r in cur.fetchall() if r[0]])
                cur.execute(
                    "SELECT face_crop_mongo_id FROM recognition_logs WHERE person_id = %s;",
                    (person_id,),
                )
                mongo_ids.extend([r[0] for r in cur.fetchall() if r[0]])

                # Xóa theo thứ tự FK
                cur.execute("DELETE FROM recognition_logs WHERE person_id = %s;", (person_id,))
                cur.execute("DELETE FROM visits WHERE person_id = %s;", (person_id,))
                cur.execute("DELETE FROM face_embeddings WHERE person_id = %s;", (person_id,))
                cur.execute("DELETE FROM persons WHERE id = %s;", (person_id,))

                conn.commit()
                logger.info(f"[PG] Deleted person: {person_id}")
                return True, list(set(mongo_ids))
        except Exception as e:
            conn.rollback()
            logger.error(f"[PG] delete_person failed: {e}")
            raise
        finally:
            conn.close()

    # ── delete_person_by_name ────────────────────────────────────────────

    def delete_person_by_name(self, name: str) -> tuple[bool, list[str]]:
        """Xóa person theo full_name. Convenience wrapper."""
        conn = self._conn()
        try:
            with conn.cursor() as cur:
                cur.execute("SELECT id FROM persons WHERE full_name = %s;", (name,))
                row = cur.fetchone()
                if not row:
                    return False, []
                person_id = str(row[0])
            return self.delete_person(person_id)
        finally:
            conn.close()

    # ── list_persons ─────────────────────────────────────────────────────

    def list_persons_with_faces(self) -> list[dict]:
        """
        Lấy danh sách tất cả persons có face_embedding, kèm mongo_id ảnh crop.
        """
        query = """
            SELECT
                p.id::text        AS person_id,
                p.full_name,
                p.status,
                p.total_appearances,
                p.total_visits,
                p.last_seen_at,
                e.id::text        AS embedding_id,
                e.face_crop_mongo_id
            FROM persons p
            JOIN face_embeddings e ON p.id = e.person_id
            WHERE e.is_active = TRUE
            ORDER BY p.full_name;
        """
        conn = self._conn()
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute(query)
                rows = cur.fetchall()
                result = []
                for r in rows:
                    d = dict(r)
                    if d.get("last_seen_at"):
                        d["last_seen_at"] = d["last_seen_at"].isoformat()
                    result.append(d)
                return result
        finally:
            conn.close()


# ── Singleton ─────────────────────────────────────────────────────────────
pg_storage = PostgresStorage()