services/face_index_service.py

"""
FaceIndexService — high-level reverse face search service.

Handles:
  - Image upload + validation
  - Face detection (using the feature extractor)
  - Embedding generation (via InsightFace provider if available)
  - Enrollment into the FaceIndex
  - Search against the FaceIndex
  - Crawler for seeding the index with public faces
"""

from __future__ import annotations

import base64
import io
import time
import uuid
from typing import Optional

import numpy as np
from loguru import logger

from config.settings import Settings
from pipeline import InputValidator, ImagePreprocessor, ImageHasher, FeatureExtractor
from storage.face_index import FaceIndex
from utils.logging import execution_context, new_execution_id


class FaceIndexService:
    """High-level service for the reverse face search index."""

    def __init__(
        self,
        settings: Settings,
        validator: InputValidator,
        preprocessor: ImagePreprocessor,
        hasher: ImageHasher,
        feature_extractor: FeatureExtractor,
        face_index: FaceIndex,
    ) -> None:
        self._settings = settings
        self._validator = validator
        self._preprocessor = preprocessor
        self._hasher = hasher
        self._feature_extractor = feature_extractor
        self._index = face_index

    # ------------------------------------------------------------------ #
    # Enrollment
    # ------------------------------------------------------------------ #
    async def enroll(
        self,
        image_url: Optional[str] = None,
        image_base64: Optional[str] = None,
        image_bytes: Optional[bytes] = None,
        name: Optional[str] = None,
        source_url: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> dict:
        """
        Enroll a face into the searchable index.

        Accepts an image (URL, base64, or raw bytes) + optional metadata.
        Detects the largest face in the image, computes its embedding,
        and adds it to the FaceIndex.

        Returns dict with face_id + status.
        """
        eid = new_execution_id()
        with execution_context(execution_id=eid, provider_id="face_index_service.enroll"):
            t0 = time.perf_counter()

            # Validate input
            if image_bytes:
                vr = self._validator.validate(image_bytes=image_bytes)
            else:
                vr = self._validator.validate(
                    image_url=image_url, image_base64=image_base64,
                )
            if not vr.valid:
                return {"success": False, "error": vr.error, "error_type": "ValidationError"}

            # Preprocess
            if vr.source == "url":
                pre = self._preprocessor.from_url(image_url)
            else:
                pre = self._preprocessor.from_bytes(vr.image_bytes, vr.source)

            img_hash = self._hasher.hash(pre.image)
            pipeline_output = self._feature_extractor.extract(
                pre.image, img_hash, pre.width, pre.height, pre.source,
                original_bytes=pre.original_bytes,
                original_format=pre.original_format,
            )

            # Check if any faces were detected
            if not pipeline_output.face_crops:
                return {
                    "success": False,
                    "error": "No faces detected in image",
                    "error_type": "NoFaceDetected",
                }

            # Use the largest face (highest area)
            largest = max(
                pipeline_output.face_crops,
                key=lambda c: c.width * c.height if hasattr(c, "width") else 0,
            )
            embedding = getattr(largest, "embedding", None)
            if embedding is None:
                # Fallback: try to compute embedding via feature extractor
                embedding = self._feature_extractor.compute_embedding(largest.image)
                if embedding is None:
                    return {
                        "success": False,
                        "error": "Could not compute face embedding (InsightFace provider not available?)",
                        "error_type": "EmbeddingFailed",
                    }

            # Enroll into index
            face_id = self._index.enroll(
                embedding=embedding,
                name=name,
                source_url=source_url,
                metadata=metadata,
                thumbnail_path=None,  # could save crop here in future
            )

            elapsed = (time.perf_counter() - t0) * 1000.0
            return {
                "success": True,
                "face_id": face_id,
                "embedding_dim": self._index.EMBEDDING_DIM,
                "elapsed_ms": round(elapsed, 3),
            }

    # ------------------------------------------------------------------ #
    # Search
    # ------------------------------------------------------------------ #
    async def search(
        self,
        image_url: Optional[str] = None,
        image_base64: Optional[str] = None,
        image_bytes: Optional[bytes] = None,
        top_k: Optional[int] = None,
        threshold: Optional[float] = None,
    ) -> dict:
        """
        Search the index for faces matching the query image.

        Returns dict with matches sorted by similarity (descending).
        """
        eid = new_execution_id()
        with execution_context(execution_id=eid, provider_id="face_index_service.search"):
            t0 = time.perf_counter()

            # Validate input
            if image_bytes:
                vr = self._validator.validate(image_bytes=image_bytes)
            else:
                vr = self._validator.validate(
                    image_url=image_url, image_base64=image_base64,
                )
            if not vr.valid:
                return {"success": False, "error": vr.error, "error_type": "ValidationError"}

            # Preprocess
            if vr.source == "url":
                pre = self._preprocessor.from_url(image_url)
            else:
                pre = self._preprocessor.from_bytes(vr.image_bytes, vr.source)

            img_hash = self._hasher.hash(pre.image)
            pipeline_output = self._feature_extractor.extract(
                pre.image, img_hash, pre.width, pre.height, pre.source,
                original_bytes=pre.original_bytes,
                original_format=pre.original_format,
            )

            if not pipeline_output.face_crops:
                return {
                    "success": True,
                    "query_face_detected": False,
                    "matches": [],
                    "elapsed_ms": round((time.perf_counter() - t0) * 1000, 3),
                }

            # Use the largest face
            largest = max(
                pipeline_output.face_crops,
                key=lambda c: c.width * c.height if hasattr(c, "width") else 0,
            )
            embedding = getattr(largest, "embedding", None)
            if embedding is None:
                embedding = self._feature_extractor.compute_embedding(largest.image)
                if embedding is None:
                    return {
                        "success": False,
                        "error": "Could not compute face embedding",
                        "error_type": "EmbeddingFailed",
                    }

            # Search the index
            matches = self._index.search(
                embedding,
                top_k=top_k or self._settings.face_index_top_k,
                threshold=threshold if threshold is not None else self._settings.face_index_threshold,
            )

            elapsed = (time.perf_counter() - t0) * 1000.0
            return {
                "success": True,
                "query_face_detected": True,
                "num_matches": len(matches),
                "matches": matches,
                "elapsed_ms": round(elapsed, 3),
            }

    # ------------------------------------------------------------------ #
    # CRUD passthrough
    # ------------------------------------------------------------------ #
    def list(self, limit: int = 50, offset: int = 0, name: Optional[str] = None) -> list[dict]:
        return self._index.list(limit=limit, offset=offset, name=name)

    def get(self, face_id: str) -> Optional[dict]:
        return self._index.get(face_id)

    def delete(self, face_id: str) -> bool:
        return self._index.delete(face_id)

    def stats(self) -> dict:
        return self._index.stats()

    def clear(self) -> int:
        return self._index.clear()