storage/face_index.py

"""
FaceIndex — sqlite-vec backed reverse face search index.

Stores 512-d ArcFace face embeddings in a SQLite database with sqlite-vec
for fast cosine similarity search.  Persists to a single file
(data/face_index.db by default).

Why sqlite-vec instead of ChromaDB / FAISS / Pinecone?
  - Zero external dependencies beyond the `sqlite-vec` Python package (5MB)
  - Persists to a single file (easy backup, easy ship)
  - Runs on free-tier VPS without RAM issues
  - Supports standard SQL queries alongside vector search
  - Can be inspected with the `sqlite3` CLI

Schema:
  faces          — one row per enrolled face (face_id, embedding, metadata)
  enrollments    — audit log of all enrollment events

The embedding is stored as a serialized 512-d float32 vector via sqlite-vec's
vec0 virtual table type.  Search is cosine similarity (since embeddings are
L2-normalized, this equals dot product).
"""

from __future__ import annotations

import json
import sqlite3
import threading
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

import numpy as np
from loguru import logger


# Schema for the metadata table (regular SQLite)
_SCHEMA_METADATA = """
CREATE TABLE IF NOT EXISTS faces (
    face_id         TEXT PRIMARY KEY,
    name            TEXT,
    source_url      TEXT,
    metadata        TEXT NOT NULL DEFAULT '{}',
    thumbnail_path  TEXT,
    embedding_dim   INTEGER NOT NULL,
    created_at      TEXT NOT NULL,
    updated_at      TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_faces_name ON faces(name);
CREATE INDEX IF NOT EXISTS idx_faces_created ON faces(created_at);

CREATE TABLE IF NOT EXISTS enrollments (
    id           INTEGER PRIMARY KEY AUTOINCREMENT,
    face_id      TEXT NOT NULL,
    action       TEXT NOT NULL,  -- 'enroll' | 'delete'
    timestamp    TEXT NOT NULL,
    details      TEXT,
    FOREIGN KEY (face_id) REFERENCES faces(face_id)
);

CREATE INDEX IF NOT EXISTS idx_enrollments_face ON enrollments(face_id);
"""


class FaceIndex:
    """Thread-safe reverse face search index backed by sqlite-vec."""

    EMBEDDING_DIM = 512

    def __init__(self, path: str = ":memory:") -> None:
        self._path = path
        self._lock = threading.RLock()

        # Ensure parent dir exists
        if path != ":memory:":
            Path(path).parent.mkdir(parents=True, exist_ok=True)

        self._conn = sqlite3.connect(path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row

        # Load sqlite-vec extension
        try:
            import sqlite_vec  # type: ignore
            self._conn.enable_load_extension(True)
            sqlite_vec.load(self._conn)
            self._conn.enable_load_extension(False)
            self._vec_available = True
            logger.info("sqlite-vec extension loaded successfully")
        except ImportError:
            self._vec_available = False
            logger.warning(
                "sqlite-vec not installed — FaceIndex will use pure-Python fallback "
                "(slower, no vector indexing). Install with: pip install sqlite-vec"
            )
        except Exception as e:
            self._vec_available = False
            logger.error(f"Failed to load sqlite-vec: {e}")

        # Initialize schema
        self._conn.executescript(_SCHEMA_METADATA)
        self._uses_l2 = False  # set by _init_vec_table
        self._init_vec_table()
        self._conn.commit()

        if self._vec_available:
            count = self.count()
            logger.info(f"FaceIndex initialized at {path} (sqlite-vec mode, {count} faces)")
        else:
            logger.info(f"FaceIndex initialized at {path} (fallback mode)")

    def _init_vec_table(self) -> None:
        """Create the vec0 virtual table for fast vector search.

        We use cosine distance (1 - dot product for L2-normalized vectors)
        by partition-adding a `distance_metric` column metadata.  However,
        sqlite-vec's vec0 currently defaults to L2 (Euclidean) distance.
        For L2-normalized embeddings, L2² = 2(1 - cos), so we can convert
        at query time: cos_sim = 1 - (L2² / 2).
        """
        if not self._vec_available:
            return
        try:
            self._conn.execute(
                f"""
                CREATE VIRTUAL TABLE IF NOT EXISTS face_embeddings USING vec0(
                    face_id TEXT PRIMARY KEY,
                    embedding float[{self.EMBEDDING_DIM}] distance_metric=cosine
                )
                """
            )
            self._conn.commit()
        except Exception as e:
            # Fallback: try without distance_metric (older sqlite-vec versions)
            logger.warning(f"cosine metric not supported, falling back to L2: {e}")
            try:
                self._conn.execute(
                    f"""
                    CREATE VIRTUAL TABLE IF NOT EXISTS face_embeddings USING vec0(
                        face_id TEXT PRIMARY KEY,
                        embedding float[{self.EMBEDDING_DIM}]
                    )
                    """
                )
                self._conn.commit()
                self._uses_l2 = True
            except Exception as e2:
                logger.error(f"Failed to create vec0 table: {e2}")
                self._vec_available = False
        else:
            self._uses_l2 = False

    # ------------------------------------------------------------------ #
    # Enrollment
    # ------------------------------------------------------------------ #
    def enroll(
        self,
        embedding: np.ndarray,
        name: Optional[str] = None,
        source_url: Optional[str] = None,
        metadata: Optional[dict] = None,
        thumbnail_path: Optional[str] = None,
    ) -> str:
        """
        Add a face to the searchable index.

        Args:
            embedding: 512-d L2-normalized face embedding (ArcFace)
            name: optional human-readable name
            source_url: optional URL where the face was found
            metadata: optional dict of arbitrary metadata (age, location, etc.)
            thumbnail_path: optional path to a thumbnail image of the face

        Returns:
            face_id (UUID string)
        """
        if embedding.shape != (self.EMBEDDING_DIM,):
            raise ValueError(
                f"Embedding must be shape ({self.EMBEDDING_DIM},), got {embedding.shape}"
            )

        # L2-normalize (defensive — should already be normalized)
        norm = np.linalg.norm(embedding)
        if norm > 0:
            embedding = embedding / norm

        face_id = str(uuid.uuid4())
        now = datetime.now(timezone.utc).isoformat()
        metadata_json = json.dumps(metadata or {}, default=str)
        embedding_bytes = embedding.astype(np.float32).tobytes()

        with self._lock:
            # Insert metadata
            self._conn.execute(
                """INSERT INTO faces
                   (face_id, name, source_url, metadata, thumbnail_path,
                    embedding_dim, created_at, updated_at)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                (face_id, name, source_url, metadata_json, thumbnail_path,
                 self.EMBEDDING_DIM, now, now),
            )

            # Insert embedding
            if self._vec_available:
                self._conn.execute(
                    "INSERT INTO face_embeddings (face_id, embedding) VALUES (?, ?)",
                    (face_id, embedding_bytes),
                )

            # Audit log
            self._conn.execute(
                """INSERT INTO enrollments (face_id, action, timestamp, details)
                   VALUES (?, ?, ?, ?)""",
                (face_id, "enroll", now, json.dumps({"name": name, "source_url": source_url})),
            )
            self._conn.commit()

        logger.debug(f"Enrolled face {face_id} (name={name})")
        return face_id

    # ------------------------------------------------------------------ #
    # Search
    # ------------------------------------------------------------------ #
    def search(
        self,
        query_embedding: np.ndarray,
        top_k: int = 10,
        threshold: float = 0.0,
    ) -> list[dict]:
        """
        Search the index for faces similar to the query embedding.

        Args:
            query_embedding: 512-d L2-normalized face embedding
            top_k: number of results to return
            threshold: minimum cosine similarity (0-1); 0 = return all

        Returns:
            list of dicts sorted by similarity (descending):
              {
                "face_id": str,
                "name": str | None,
                "source_url": str | None,
                "metadata": dict,
                "thumbnail_path": str | None,
                "similarity": float,
                "created_at": str,
              }
        """
        if query_embedding.shape != (self.EMBEDDING_DIM,):
            raise ValueError(
                f"Query embedding must be shape ({self.EMBEDDING_DIM},), got {query_embedding.shape}"
            )

        # L2-normalize
        norm = np.linalg.norm(query_embedding)
        if norm > 0:
            query_embedding = query_embedding / norm

        with self._lock:
            if self._vec_available:
                matches = self._search_vec(query_embedding, top_k * 2)
            else:
                matches = self._search_fallback(query_embedding, top_k * 2)

        # Filter by threshold + take top_k
        results = []
        for m in matches:
            if m["similarity"] >= threshold:
                results.append(m)
            if len(results) >= top_k:
                break
        return results

    def _search_vec(self, query_embedding: np.ndarray, k: int) -> list[dict]:
        """Use sqlite-vec KNN search."""
        query_bytes = query_embedding.astype(np.float32).tobytes()
        rows = self._conn.execute(
            """
            SELECT
                f.face_id, f.name, f.source_url, f.metadata,
                f.thumbnail_path, f.created_at,
                v.distance
            FROM face_embeddings v
            JOIN faces f ON f.face_id = v.face_id
            WHERE v.embedding MATCH ?
              AND k = ?
            ORDER BY v.distance ASC
            """,
            (query_bytes, k),
        ).fetchall()

        # distance interpretation depends on metric:
        # - cosine: distance = 1 - cos_sim, so sim = 1 - distance
        # - L2 (Euclidean): for L2-normalized vectors, L2² = 2(1 - cos_sim),
        #   so cos_sim = 1 - (distance² / 2).  We use squared distance here
        #   since sqlite-vec returns Euclidean (not squared).
        out = []
        for r in rows:
            if getattr(self, "_uses_l2", False):
                # L2 distance — convert to cosine similarity
                sim = 1.0 - (r["distance"] ** 2) / 2.0
            else:
                # cosine distance
                sim = 1.0 - r["distance"]
            out.append({
                "face_id": r["face_id"],
                "name": r["name"],
                "source_url": r["source_url"],
                "metadata": json.loads(r["metadata"] or "{}"),
                "thumbnail_path": r["thumbnail_path"],
                "similarity": float(sim),
                "created_at": r["created_at"],
            })
        return out

    def _search_fallback(self, query_embedding: np.ndarray, k: int) -> list[dict]:
        """Pure-Python fallback when sqlite-vec is not available."""
        rows = self._conn.execute(
            "SELECT face_id, name, source_url, metadata, thumbnail_path, created_at FROM faces"
        ).fetchall()

        # We don't store embeddings in the metadata table for fallback mode.
        # In a real fallback, we'd need a separate embedding store. For now,
        # return empty results — installation of sqlite-vec is required.
        logger.warning(
            "FaceIndex fallback mode does not support search — install sqlite-vec: "
            "pip install sqlite-vec"
        )
        return []

    # ------------------------------------------------------------------ #
    # CRUD
    # ------------------------------------------------------------------ #
    def get(self, face_id: str) -> Optional[dict]:
        """Get details of a specific enrolled face."""
        with self._lock:
            row = self._conn.execute(
                "SELECT * FROM faces WHERE face_id = ?", (face_id,)
            ).fetchone()
        if not row:
            return None
        d = dict(row)
        d["metadata"] = json.loads(d.get("metadata") or "{}")
        return d

    def list(self, limit: int = 50, offset: int = 0, name: Optional[str] = None) -> list[dict]:
        """Paginated list of enrolled faces."""
        with self._lock:
            if name:
                cur = self._conn.execute(
                    "SELECT * FROM faces WHERE name LIKE ? ORDER BY created_at DESC LIMIT ? OFFSET ?",
                    (f"%{name}%", limit, offset),
                )
            else:
                cur = self._conn.execute(
                    "SELECT * FROM faces ORDER BY created_at DESC LIMIT ? OFFSET ?",
                    (limit, offset),
                )
            rows = [dict(r) for r in cur.fetchall()]
        for r in rows:
            r["metadata"] = json.loads(r.get("metadata") or "{}")
        return rows

    def delete(self, face_id: str) -> bool:
        """Remove a face from the index. Returns True if deleted."""
        now = datetime.now(timezone.utc).isoformat()
        with self._lock:
            # Check exists
            row = self._conn.execute(
                "SELECT face_id FROM faces WHERE face_id = ?", (face_id,)
            ).fetchone()
            if not row:
                return False

            self._conn.execute("DELETE FROM faces WHERE face_id = ?", (face_id,))
            if self._vec_available:
                self._conn.execute(
                    "DELETE FROM face_embeddings WHERE face_id = ?", (face_id,)
                )
            self._conn.execute(
                """INSERT INTO enrollments (face_id, action, timestamp, details)
                   VALUES (?, ?, ?, ?)""",
                (face_id, "delete", now, "{}"),
            )
            self._conn.commit()
        logger.debug(f"Deleted face {face_id}")
        return True

    def count(self) -> int:
        """Total number of enrolled faces."""
        with self._lock:
            row = self._conn.execute("SELECT COUNT(*) as n FROM faces").fetchone()
        return row["n"] if row else 0

    def stats(self) -> dict:
        """Index statistics."""
        with self._lock:
            total = self.count()
            named = self._conn.execute(
                "SELECT COUNT(*) as n FROM faces WHERE name IS NOT NULL"
            ).fetchone()["n"]
            last_enrollment = self._conn.execute(
                "SELECT created_at FROM faces ORDER BY created_at DESC LIMIT 1"
            ).fetchone()
            recent_enrollments = self._conn.execute(
                """SELECT COUNT(*) as n FROM enrollments
                   WHERE action = 'enroll'
                     AND timestamp > datetime('now', '-24 hours')"""
            ).fetchone()["n"]
        return {
            "total_faces": total,
            "named_faces": named,
            "anonymous_faces": total - named,
            "last_enrollment": last_enrollment["created_at"] if last_enrollment else None,
            "recent_enrollments_24h": recent_enrollments,
            "vec_available": self._vec_available,
            "embedding_dim": self.EMBEDDING_DIM,
            "path": self._path,
        }

    def clear(self) -> int:
        """Remove ALL faces from the index. Returns count deleted."""
        with self._lock:
            n = self.count()
            self._conn.execute("DELETE FROM faces")
            if self._vec_available:
                self._conn.execute("DELETE FROM face_embeddings")
            self._conn.execute("DELETE FROM enrollments")
            self._conn.commit()
        logger.warning(f"Cleared {n} faces from index")
        return n

    def close(self) -> None:
        with self._lock:
            self._conn.close()