src/agents/hf_embedding.py

"""
Embedding agent with HuggingFace API or local fallback.

Uses HuggingFace Inference API when API key is available,
otherwise falls back to local sentence-transformers.
"""

from __future__ import annotations

import logging
import time
from typing import Callable

import numpy as np

from src.agents.base import Agent, AgentContext
from src.config import HF_API_KEY

logger = logging.getLogger(__name__)

# Model name mapping (short name -> HF model ID)
EMBEDDING_MODELS = {
    "all-MiniLM-L6-v2": "sentence-transformers/all-MiniLM-L6-v2",
    "all-mpnet-base-v2": "sentence-transformers/all-mpnet-base-v2",
    "bge-large": "BAAI/bge-large-en-v1.5",
    "bge-base": "BAAI/bge-base-en-v1.5",
    "bge-small": "BAAI/bge-small-en-v1.5",
}


class EmbeddingClient:
    """
    Client for getting embeddings via HuggingFace API or local models.

    Automatically uses HF API if API key is available, otherwise local.
    """

    def __init__(
        self,
        model_name: str = "all-MiniLM-L6-v2",
        api_key: str | None = None,
        on_status: Callable[[str], None] | None = None,
        force_local: bool = False,
    ) -> None:
        """
        Initialize the embedding client.

        Args:
            model_name: Short name or full HF model ID
            api_key: HuggingFace API key (uses env var if not provided)
            on_status: Callback for status messages (warmup, loading, etc.)
            force_local: Force local model even if API key is available
        """
        # Resolve model name
        if model_name in EMBEDDING_MODELS:
            self._model_id = EMBEDDING_MODELS[model_name]
        else:
            self._model_id = model_name

        self._short_name = model_name.split("/")[-1]
        self._api_key = api_key or HF_API_KEY
        self._on_status = on_status
        self._force_local = force_local
        self._is_warm = False

        # Will be initialized lazily
        self._hf_client = None
        self._local_model = None
        self._use_api = bool(self._api_key) and not force_local

    @property
    def model_id(self) -> str:
        return self._model_id

    @property
    def is_warm(self) -> bool:
        return self._is_warm

    @property
    def using_api(self) -> bool:
        return self._use_api

    def _notify(self, message: str) -> None:
        """Send status notification."""
        if self._on_status:
            self._on_status(message)
        logger.info(message)

    def _init_hf_client(self) -> bool:
        """Initialize HuggingFace API client."""
        if self._hf_client is not None:
            return True

        try:
            from huggingface_hub import InferenceClient
            self._hf_client = InferenceClient(token=self._api_key)
            return True
        except ImportError:
            logger.warning("huggingface_hub not installed, falling back to local")
            self._use_api = False
            return False
        except Exception as e:
            logger.warning(f"Failed to init HF client: {e}, falling back to local")
            self._use_api = False
            return False

    def _init_local_model(self) -> bool:
        """Initialize local sentence-transformer model."""
        if self._local_model is not None:
            return True

        try:
            self._notify(f"Loading model locally ({self._short_name})...")
            from sentence_transformers import SentenceTransformer
            self._local_model = SentenceTransformer(self._model_id)
            self._notify("Model loaded!")
            return True
        except Exception as e:
            logger.error(f"Failed to load local model: {e}")
            return False

    def _embed_via_api(self, texts: list[str]) -> np.ndarray | None:
        """Get embeddings via HuggingFace API."""
        if not self._init_hf_client():
            return None

        try:
            self._notify("Getting embeddings from HuggingFace API...")

            # HF API handles batching internally
            embeddings = []
            for text in texts:
                result = self._hf_client.feature_extraction(
                    text,
                    model=self._model_id,
                )
                # Result may be nested, flatten if needed
                if isinstance(result, list) and len(result) > 0:
                    if isinstance(result[0], list):
                        # Token-level embeddings, take mean
                        embeddings.append(np.mean(result, axis=0))
                    else:
                        embeddings.append(result)
                else:
                    embeddings.append(result)

            self._is_warm = True
            return np.array(embeddings)

        except Exception as e:
            error_str = str(e)
            if "503" in error_str or "loading" in error_str.lower():
                self._notify("Model warming up on HuggingFace (~20s)...")
                time.sleep(20)
                return self._embed_via_api(texts)  # Retry
            elif "401" in error_str or "unauthorized" in error_str.lower():
                logger.warning("HF API unauthorized, falling back to local")
                self._use_api = False
                return self._embed_via_local(texts)
            else:
                logger.error(f"HF API error: {e}")
                # Fall back to local
                self._use_api = False
                return self._embed_via_local(texts)

    def _embed_via_local(self, texts: list[str]) -> np.ndarray | None:
        """Get embeddings via local model."""
        if not self._init_local_model():
            return None

        try:
            embeddings = self._local_model.encode(texts, convert_to_numpy=True)
            self._is_warm = True
            return embeddings
        except Exception as e:
            logger.error(f"Local embedding error: {e}")
            return None

    def embed(self, texts: list[str]) -> np.ndarray | None:
        """
        Get embeddings for a list of texts.

        Automatically uses HF API if available, otherwise local model.

        Args:
            texts: Texts to embed

        Returns:
            Numpy array of embeddings (n_texts, embedding_dim) or None
        """
        if not texts:
            return np.array([])

        if self._use_api:
            return self._embed_via_api(texts)
        else:
            return self._embed_via_local(texts)

    def warmup(self) -> bool:
        """
        Warm up the model with a test query.

        Returns:
            True if model is ready, False otherwise
        """
        try:
            self._notify("Warming up embedding model...")
            result = self.embed(["test warmup query"])
            if result is not None and len(result) > 0:
                self._notify("Model ready!")
                return True
            return False
        except Exception as e:
            logger.error(f"Warmup failed: {e}")
            return False


class HFEmbeddingAgent(Agent):
    """
    Greedy agent that uses embeddings for link selection.

    Automatically uses HuggingFace API when available, otherwise local models.
    """

    def __init__(
        self,
        model_name: str = "all-MiniLM-L6-v2",
        avoid_revisits: bool = True,
        on_status: Callable[[str], None] | None = None,
        force_local: bool = False,
    ) -> None:
        """
        Initialize embedding agent.

        Args:
            model_name: Model name (short name or full HF model ID)
            avoid_revisits: If True, penalize revisiting pages
            on_status: Callback for status messages
            force_local: Force local model even if API key available
        """
        self._model_name = model_name
        self._avoid_revisits = avoid_revisits
        self._client = EmbeddingClient(
            model_name,
            on_status=on_status,
            force_local=force_local,
        )

    @property
    def name(self) -> str:
        short_name = self._model_name.split("/")[-1]
        prefix = "hf" if self._client.using_api else "local"
        return f"{prefix}-{short_name}"

    @property
    def description(self) -> str:
        mode = "HuggingFace API" if self._client.using_api else "local"
        return f"Greedy embedding ({mode}: {self._model_name})"

    def warmup(self) -> bool:
        """Warm up the model before gameplay."""
        return self._client.warmup()

    def _compute_similarities(
        self, candidates: list[str], target: str
    ) -> list[tuple[str, float]]:
        """Compute cosine similarities between candidates and target."""
        if not candidates:
            return []

        # Encode target and all candidates
        all_texts = [target] + candidates
        embeddings = self._client.embed(all_texts)

        if embeddings is None or len(embeddings) == 0:
            logger.warning("Failed to get embeddings")
            return [(c, 0.0) for c in candidates]

        # Normalize for cosine similarity
        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
        norms = np.where(norms == 0, 1, norms)  # Avoid division by zero
        embeddings = embeddings / norms

        # Target is first embedding
        target_emb = embeddings[0]
        candidate_embs = embeddings[1:]

        # Compute similarities
        similarities = np.dot(candidate_embs, target_emb)

        # Return sorted by similarity
        results = list(zip(candidates, similarities.tolist(), strict=True))
        results.sort(key=lambda x: x[1], reverse=True)
        return [(title, sim) for title, sim in results]

    def choose_link(self, context: AgentContext) -> str:
        """Pick the link most similar to target using embeddings."""
        # Always click target if available
        if context.target_title in context.available_links:
            return context.target_title

        # Compute similarities
        ranked = self._compute_similarities(
            candidates=context.available_links,
            target=context.target_title,
        )

        if not ranked:
            return context.available_links[0]

        # Filter out revisits if enabled
        if self._avoid_revisits:
            visited = set(context.path_so_far)
            for title, _sim in ranked:
                if title not in visited:
                    return title

        return ranked[0][0]