src/db.py

"""
src/db.py — All external service I/O lives here.

Two responsibilities:
  1. PineconePool  — LRU cache of Pinecone client objects (avoids per-request TCP overhead)
  2. Cloudinary helpers — thin wrappers that always inject credentials explicitly
                          (stateless; no global cloudinary.config() call needed)

Search helpers (search_faces, search_objects) live here too so main.py
stays thin and the query logic is testable in isolation.
"""

import io
from collections import OrderedDict
from typing import Optional

import cloudinary
import cloudinary.api
import cloudinary.uploader
from pinecone import Pinecone, ServerlessSpec

from .config import (
    PINECONE_POOL_MAX,
    IDX_FACES, IDX_OBJECTS,
    IDX_FACES_DIM, IDX_OBJECTS_DIM,
    FACE_THRESHOLD_HIGH, FACE_THRESHOLD_LOW, FACE_DET_SCORE_HQ_SPLIT,
    FACE_TOP_K_FETCH,
    OBJECT_SCORE_THRESHOLD, OBJECT_TOP_K,
)
from .utils import to_list, face_ui_score


# ════════════════════════════════════════════════════════════════════
#  PINECONE LRU CONNECTION POOL
# ════════════════════════════════════════════════════════════════════

class PineconePool:
    """
    LRU cache of Pinecone() client objects keyed by API key.

    Why pool?  Creating a Pinecone() client opens a TCP connection and
    authenticates — adds ~200 ms on first use.  Re-using clients across
    requests eliminates that overhead for repeat users.

    Implementation:  OrderedDict + move_to_end = O(1) LRU without any
    extra dependency.  On overflow, popitem(last=False) evicts the
    least-recently-used entry.
    """

    def __init__(self, maxsize: int = PINECONE_POOL_MAX):
        self._pool: OrderedDict[str, Pinecone] = OrderedDict()
        self._maxsize = maxsize

    def get(self, api_key: str) -> Pinecone:
        if api_key not in self._pool:
            if len(self._pool) >= self._maxsize:
                self._pool.popitem(last=False)
            self._pool[api_key] = Pinecone(api_key=api_key)
        self._pool.move_to_end(api_key)
        return self._pool[api_key]


# Module-level singleton — shared across all requests
pinecone_pool = PineconePool()


# ════════════════════════════════════════════════════════════════════
#  CLOUDINARY WRAPPERS
#  All functions accept a `creds` dict — no global state.
#  Call via asyncio.to_thread() from async endpoints.
# ════════════════════════════════════════════════════════════════════

def cld_upload(file_obj: io.BytesIO, folder: str, creds: dict) -> dict:
    """Upload an in-memory file to Cloudinary. Returns the full API response."""
    return cloudinary.uploader.upload(
        file_obj, folder=folder,
        api_key=creds["api_key"], api_secret=creds["api_secret"],
        cloud_name=creds["cloud_name"],
    )


def cld_ping(creds: dict) -> dict:
    """Lightweight connectivity check — verifies credentials are valid."""
    return cloudinary.api.ping(
        api_key=creds["api_key"], api_secret=creds["api_secret"],
        cloud_name=creds["cloud_name"],
    )


def cld_root_folders(creds: dict) -> dict:
    """List all top-level folders in the account."""
    return cloudinary.api.root_folders(
        api_key=creds["api_key"], api_secret=creds["api_secret"],
        cloud_name=creds["cloud_name"],
    )


def cld_list_folder_images(
    folder: str,
    creds: dict,
    next_cursor: Optional[str] = None,
    max_results: int = 100,
) -> dict:
    """Paginated listing of resources under a folder prefix."""
    kwargs = dict(
        type="upload",
        prefix=f"{folder}/",
        max_results=min(max_results, 100),   # Cloudinary hard-caps at 100
        api_key=creds["api_key"],
        api_secret=creds["api_secret"],
        cloud_name=creds["cloud_name"],
    )
    if next_cursor:
        kwargs["next_cursor"] = next_cursor
    return cloudinary.api.resources(**kwargs)


def cld_delete_resource(public_id: str, creds: dict) -> dict:
    """Delete a single resource by public_id."""
    return cloudinary.uploader.destroy(
        public_id,
        api_key=creds["api_key"], api_secret=creds["api_secret"],
        cloud_name=creds["cloud_name"],
    )


def cld_delete_folder_resources(folder: str, creds: dict) -> dict:
    """Delete all resources under a folder prefix (not the folder record itself)."""
    return cloudinary.api.delete_resources_by_prefix(
        f"{folder}/",
        api_key=creds["api_key"], api_secret=creds["api_secret"],
        cloud_name=creds["cloud_name"],
    )


def cld_remove_folder(folder: str, creds: dict):
    """Delete the folder record itself. Silently ignores errors (already gone)."""
    try:
        return cloudinary.api.delete_folder(
            folder,
            api_key=creds["api_key"], api_secret=creds["api_secret"],
            cloud_name=creds["cloud_name"],
        )
    except Exception:
        pass


def cld_delete_all_paginated(creds: dict) -> int:
    """
    Delete ALL resources in a Cloudinary account, paginating in batches of 100.
    Returns the total number of resources deleted.
    Used only by the destructive /api/reset-database and /api/delete-account endpoints.
    """
    deleted = 0
    while True:
        try:
            res = cloudinary.api.resources(
                type="upload", max_results=100,
                api_key=creds["api_key"], api_secret=creds["api_secret"],
                cloud_name=creds["cloud_name"],
            )
            resources = res.get("resources", [])
            if not resources:
                break
            public_ids = [r["public_id"] for r in resources]
            cloudinary.api.delete_resources(
                public_ids,
                api_key=creds["api_key"], api_secret=creds["api_secret"],
                cloud_name=creds["cloud_name"],
            )
            deleted += len(public_ids)
            if not res.get("next_cursor"):
                break
        except Exception:
            break
    return deleted


# ════════════════════════════════════════════════════════════════════
#  PINECONE INDEX MANAGEMENT
# ════════════════════════════════════════════════════════════════════

def ensure_indexes(pc: Pinecone) -> list[str]:
    """
    Create enterprise-faces and enterprise-objects indexes if they don't exist.
    Returns a list of index names that were freshly created (empty list if both existed).
    """
    existing       = {idx.name for idx in pc.list_indexes()}
    created        = []
    spec           = ServerlessSpec(cloud="aws", region="us-east-1")

    if IDX_OBJECTS not in existing:
        pc.create_index(name=IDX_OBJECTS, dimension=IDX_OBJECTS_DIM, metric="cosine", spec=spec)
        created.append(IDX_OBJECTS)

    if IDX_FACES not in existing:
        pc.create_index(name=IDX_FACES, dimension=IDX_FACES_DIM, metric="cosine", spec=spec)
        created.append(IDX_FACES)

    return created


def delete_and_recreate_indexes(pc: Pinecone):
    """
    Destroy both indexes and recreate them empty.
    Used by /api/reset-database. Caller is responsible for ensuring
    this is not called on the shared demo database.
    """
    import asyncio, time
    existing = {idx.name for idx in pc.list_indexes()}
    for name in [IDX_OBJECTS, IDX_FACES]:
        if name in existing:
            pc.delete_index(name)
    time.sleep(3)   # Pinecone needs a moment to fully delete before recreating
    spec = ServerlessSpec(cloud="aws", region="us-east-1")
    pc.create_index(name=IDX_OBJECTS, dimension=IDX_OBJECTS_DIM, metric="cosine", spec=spec)
    pc.create_index(name=IDX_FACES,   dimension=IDX_FACES_DIM,   metric="cosine", spec=spec)


# ════════════════════════════════════════════════════════════════════
#  PINECONE SEARCH HELPERS
# ════════════════════════════════════════════════════════════════════

def search_faces(idx_face, vec: list, det_score: float) -> dict:
    """
    Query the faces index for one face vector.

    Threshold selection:
      - High-quality face (det_score >= FACE_DET_SCORE_HQ_SPLIT) → FACE_THRESHOLD_HIGH
      - Lower-quality face                                        → FACE_THRESHOLD_LOW
    Adaptive thresholding prevents low-quality query faces from flooding results
    with false positives while still matching when confidence warrants it.

    Returns: { image_url → { raw_score, face_crop, folder, face_width_px } }
    De-duplicated per image URL — only the highest-scoring face match per image
    is retained (multiple face vectors can be stored per source image during upload).
    """
    threshold = (
        FACE_THRESHOLD_HIGH
        if det_score >= FACE_DET_SCORE_HQ_SPLIT
        else FACE_THRESHOLD_LOW
    )
    response  = idx_face.query(vector=vec, top_k=FACE_TOP_K_FETCH, include_metadata=True)
    image_map = {}

    for match in response.get("matches", []):
        raw = match["score"]
        if raw < threshold:
            continue
        url = match["metadata"].get("url", "")
        if not url:
            continue
        if url not in image_map or raw > image_map[url]["raw_score"]:
            image_map[url] = {
                "raw_score":     raw,
                "face_crop":     match["metadata"].get("face_crop",     ""),
                "folder":        match["metadata"].get("folder",        ""),
                "face_width_px": int(match["metadata"].get("face_width_px", 0)),
            }
    return image_map


def search_objects(idx_obj, vec: list) -> list:
    """
    Query the objects index for one object vector.

    Returns list of { url, score, folder } dicts for matches above
    OBJECT_SCORE_THRESHOLD, sorted by score descending.
    """
    response = idx_obj.query(vector=vec, top_k=OBJECT_TOP_K, include_metadata=True)
    results  = []
    for match in response.get("matches", []):
        if match["score"] < OBJECT_SCORE_THRESHOLD:
            continue
        url = match["metadata"].get("url", "")
        if url:
            results.append({
                "url":    url,
                "score":  round(match["score"], 4),
                "folder": match["metadata"].get("folder", ""),
            })
    return sorted(results, key=lambda x: x["score"], reverse=True)


def merge_face_results(raw_groups: list[dict]) -> list[dict]:
    """
    Merge per-face search results from multiple query faces into one ranked list.

    Algorithm:
      1. Build a global image_url → {best_score, matched_face_indices, ...} map.
      2. An image is included if ANY query face matched it.
      3. Score = best raw score across all matching faces, remapped to UI score.
      4. Multi-face boost: +5 % per additional matched face (capped at 0.99).
         Rationale: an image containing 3 of your 3 searched people should rank
         above one containing only 1 of 3.

    Args:
        raw_groups: list of dicts, each containing '_image_map' from search_faces()

    Returns: sorted list of merged result dicts (best first), max 20 items.
    """
    global_map: dict[str, dict] = {}

    for gi, group in enumerate(raw_groups):
        for url, d in group["_image_map"].items():
            raw = d["raw_score"]
            if url not in global_map:
                global_map[url] = {
                    "raw_score":     raw,
                    "face_crop":     d["face_crop"],
                    "folder":        d["folder"],
                    "face_width_px": d["face_width_px"],
                    "matched_faces": [gi],
                }
            else:
                entry = global_map[url]
                entry["matched_faces"].append(gi)
                if raw > entry["raw_score"]:
                    entry["raw_score"]     = raw
                    entry["face_crop"]     = d["face_crop"]
                    entry["face_width_px"] = d["face_width_px"]

    results = []
    for url, d in global_map.items():
        n = len(d["matched_faces"])
        results.append({
            "url":           url,
            "score":         face_ui_score(d["raw_score"], n),
            "raw_score":     round(d["raw_score"], 4),
            "face_crop":     d["face_crop"],
            "folder":        d["folder"],
            "face_width_px": d["face_width_px"],
            "matched_faces": d["matched_faces"],
            "caption":       f"👥 {n} faces matched" if n > 1 else "👤 Verified Identity",
        })

    return sorted(results, key=lambda x: x["score"], reverse=True)[:20]


def merge_object_results(nested: list[list[dict]]) -> list[dict]:
    """
    Merge object search results from multiple crops into one ranked list.

    When YOLO produces N crops, each gets its own embedding and Pinecone query.
    The same image URL may appear in multiple crop results; only the highest
    score per URL is kept.

    Returns: sorted list, max OBJECT_TOP_K items.
    """
    seen: dict[str, dict] = {}
    for results in nested:
        for r in results:
            url = r["url"]
            if url and (url not in seen or r["score"] > seen[url]["score"]):
                seen[url] = r
    return sorted(seen.values(), key=lambda x: x["score"], reverse=True)[:OBJECT_TOP_K]