app/api/faiss.py

"""
faiss.py  –  FAISS management endpoints
-----------------------------------------
Endpoints
  POST /add-embedding   – add a pre-computed embedding to the stored-item index
  GET  /download-index  – download the raw faiss_stored_items.index file

Upload behaviour
  After every /add-embedding call the updated index is pushed back to
  HuggingFace Hub (FitFinder/faiss-storage) in the background so the
  remote copy always mirrors the latest local state.

Cron behaviour
  Every call to /add-embedding resets a 15-minute inactivity timer.
  When the timer fires (no new embeddings for 15 min) the service
  re-downloads faiss_stored_items.index from HuggingFace Hub and
  reloads it into the in-process FaissService instance so the two
  stores stay in sync.
"""

import asyncio
import os
from typing import List, Optional

import numpy as np
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import FileResponse
from huggingface_hub import hf_hub_download, upload_file as hf_upload_file
from pydantic import BaseModel, field_validator

from app.services.faiss_service import FAISS_DIR, INDEX_PATHS, EMBEDDING_DIM

router = APIRouter()

# ── inactivity timer state ────────────────────────────────────────────────────
_refresh_task: Optional[asyncio.Task] = None
_INACTIVITY_SECONDS = 15 * 60  # 15 minutes


async def _refresh_after_inactivity(faiss_service) -> None:
    """Sleep for 15 minutes then push the local index to HuggingFace Hub and refresh from the remote store."""
    await asyncio.sleep(_INACTIVITY_SECONDS)

    print("[FAISS] 15‑min inactivity elapsed – uploading index to HuggingFace Hub …")
    try:
        # Upload the current local index first
        await _upload_to_hf(faiss_service.index_path)
    except Exception as exc:
        print(f"[FAISS] Upload failed during inactivity refresh: {exc}")

    print("[FAISS] Downloading latest index from HuggingFace Hub …")
    try:
        hf_hub_download(
            repo_id="FitFinder/faiss-storage",
            filename="faiss_stored_items.index",
            local_dir="/data",
            local_dir_use_symlinks=False,
            force_download=True,          # always fetch the latest version
        )
        # Reload the index file into the live service instance
        import faiss as _faiss
        new_index = _faiss.read_index(INDEX_PATHS["stored_item"])
        faiss_service.index = new_index
        print("[FAISS] Index reloaded successfully after inactivity refresh.")
    except asyncio.CancelledError:
        # Timer was reset – nothing to do
        pass
    except Exception as exc:
        print(f"[FAISS] Refresh failed: {exc}")


def _reset_inactivity_timer(faiss_service) -> None:
    """Cancel any pending refresh task and start a fresh one."""
    global _refresh_task
    if _refresh_task and not _refresh_task.done():
        _refresh_task.cancel()
    _refresh_task = asyncio.create_task(
        _refresh_after_inactivity(faiss_service)
    )


async def _upload_to_hf(index_path: str) -> None:
    """Push the local index file to HuggingFace Hub in the background."""
    try:
        print("[FAISS] Uploading updated index to HuggingFace Hub …")
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(
            None,
            lambda: hf_upload_file(
                path_or_fileobj=index_path,
                path_in_repo="faiss_stored_items.index",
                repo_id="FitFinder/faiss-storage",
                repo_type="dataset",
            ),
        )
        print("[FAISS] HuggingFace Hub upload complete.")
    except Exception as exc:
        print(f"[FAISS] HuggingFace upload failed: {exc}")


# ── request / response models ─────────────────────────────────────────────────

class AddEmbeddingRequest(BaseModel):
    embedding: List[float]

    @field_validator("embedding", mode="before")
    @classmethod
    def flatten_and_check(cls, v):
        """Accept a flat or nested (batch‑size‑1) list and ensure correct dim."""
        flat = np.array(v, dtype=np.float32).ravel().tolist()
        if len(flat) != EMBEDDING_DIM:
            raise ValueError(
                f"Embedding must contain {EMBEDDING_DIM} values, got {len(flat)}"
            )
        return flat


class AddEmbeddingResponse(BaseModel):
    faiss_id: int
    message: str


# ── endpoints ─────────────────────────────────────────────────────────────────

@router.post(
    "/add-embedding",
    response_model=AddEmbeddingResponse,
    status_code=200,
    summary="Add a pre-computed embedding to the stored-item FAISS index",
)
async def add_embedding(request: Request, body: AddEmbeddingRequest):
    """
    Accepts a flat list of floats as `embedding` in the JSON body,
    adds it to the *stored_item* FAISS index, persists the index to
    disk, and (re)starts the 15-minute inactivity timer.
    """
    faiss_service = getattr(request.app.state, "faiss_service", None)
    if faiss_service is None:
        raise HTTPException(status_code=503, detail="FAISS service not available")

    try:
        embedding_np = np.array(body.embedding, dtype=np.float32)
        faiss_id = faiss_service.add_embedding(embedding_np)
        faiss_service.save_index()
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc))
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Failed to add embedding: {exc}")

    # Reset the 15-min inactivity cron
    _reset_inactivity_timer(faiss_service)

    return AddEmbeddingResponse(
        faiss_id=faiss_id,
        message=f"Embedding added at FAISS id {faiss_id}. Index saved. "
                f"Inactivity timer reset ({_INACTIVITY_SECONDS // 60} min).",
    )


@router.get(
    "/download-index",
    summary="Download the faiss_stored_items.index file",
    response_class=FileResponse,
)
async def download_index():
    """
    Streams the raw `faiss_stored_items.index` binary file that lives
    inside the Docker container at /data/faiss_stored_items.index.
    """
    index_path = INDEX_PATHS["stored_item"]
    if not os.path.exists(index_path):
        raise HTTPException(
            status_code=404,
            detail=f"Index file not found at {index_path}",
        )

    return FileResponse(
        path=index_path,
        media_type="application/octet-stream",
        filename="faiss_stored_items.index",
    )


class RemoveEmbeddingResponse(BaseModel):
    removed_faiss_id: int
    remaining_total: int
    message: str


@router.delete(
    "/remove-embedding/{faiss_id}",
    response_model=RemoveEmbeddingResponse,
    status_code=200,
    summary="Remove an embedding from the stored-item FAISS index by its ID",
)
async def remove_embedding(request: Request, faiss_id: int):
    """
    Removes the embedding at `faiss_id` from the in-memory FAISS index,
    saves the updated index to disk, and resets the 15-minute inactivity timer.

    **⚠️ ID shift warning:** FAISS flat indexes use sequential integer IDs
    (0, 1, 2, …).  Deleting ID *N* shifts every ID greater than *N* down by 1.
    Update your own database to reflect the new IDs after calling this endpoint.
    """
    faiss_service = getattr(request.app.state, "faiss_service", None)
    if faiss_service is None:
        raise HTTPException(status_code=503, detail="FAISS service not available")

    try:
        remaining_total = faiss_service.remove_embedding(faiss_id)
        faiss_service.save_index()
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc))
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Failed to remove embedding: {exc}")

    # Reset the 15-min inactivity cron
    _reset_inactivity_timer(faiss_service)

    return RemoveEmbeddingResponse(
        removed_faiss_id=faiss_id,
        remaining_total=remaining_total,
        message=f"Embedding {faiss_id} removed. {remaining_total} vector(s) remain. "
                f"⚠️ All IDs > {faiss_id} have shifted down by 1. "
                f"Inactivity timer reset ({_INACTIVITY_SECONDS // 60} min).",
    )