"""In-memory FAISS gallery assigning persistent person IDs to face embeddings.

Duplicated from ../app/identity_store.py for this standalone web demo deployable
(see face_engine.py for why). Here, one instance lives per browser session
(held in gr.State), not per process - nothing persists between visitors.
"""
from typing import List

import faiss
import numpy as np


class IdentityStore:
    """Cosine-similarity gallery over L2-normalized ArcFace embeddings.

    Uses FAISS IndexFlatIP (inner product), which equals cosine similarity
    when all stored and query vectors are L2-normalized.
    """

    def __init__(self, embedding_dim: int = 512, match_threshold: float = 0.38) -> None:
        self.embedding_dim = embedding_dim
        self.match_threshold = match_threshold
        self.index = faiss.IndexFlatIP(embedding_dim)
        self._next_person_id = 1
        self.person_ids: List[int] = []

    def match_or_register(self, normed_embedding: np.ndarray) -> int:
        """Finds the closest known identity; registers a new one if none is close enough."""
        query = normed_embedding.astype(np.float32).reshape(1, -1)

        if self.index.ntotal > 0:
            similarities, indices = self.index.search(query, 1)
            best_similarity = float(similarities[0, 0])
            best_index = int(indices[0, 0])
            if best_similarity >= self.match_threshold:
                return self.person_ids[best_index]

        person_id = self._next_person_id
        self._next_person_id += 1
        self.index.add(query)
        self.person_ids.append(person_id)
        return person_id

    def __len__(self) -> int:
        return self.index.ntotal