"""
Embedding providers for the retrieval layer.

Two things live here:

1. A *pluggable* embedding interface (``Embedder``) plus a concrete
   OpenAI-compatible remote provider.  This is the **production** path: it hits
   an HTTP endpoint and returns real dense vectors.

2. A **deterministic local hashing fallback** (``HashingEmbedder``) that needs
   no model weights, no GPU, no API key and no internet.  It projects a bag of
   character n-grams into a fixed-dimension vector via the hashing trick.  This
   lets the whole ANN pipeline (FAISS build -> ANN search -> ranking) run
   end-to-end offline, which is exactly what we want for demos and CI.

WHY NOT A LOCAL EMBEDDING MODEL?
--------------------------------
This box is a Strix Halo AP with ~104 GiB of unified RAM.  It has *tried* to
load a large local embedding/LLM model before and OOM-killed.  So we do not
import torch or try to run sentence-transformers here.  If you have the RAM and
want real embeddings without paying for an API, drop in a local model:

    class LocalEmbedder(Embedder):
        DIM = 768
        def __init__(self):
            from sentence_transformers import SentenceTransformer
            self._model = SentenceTransformer("all-MiniLM-L6-v2")
        def embed(self, texts):
            return self._model.encode(texts, normalize_embeddings=True)

and pass it to the index/query builders via the ``embedder=`` argument.  The
rest of the pipeline is agnostic to where the vectors come from.
"""

from __future__ import annotations

import hashlib
import os
from abc import ABC, abstractmethod
from typing import Sequence

import numpy as np

# The vector dimension the hashing fallback projects into.  The remote provider
# returns its own native dimension; the index is built to whatever the embedder
# reports via ``dimension``.
HASHING_DIM = 256


class Embedder(ABC):
    """Abstract embedding provider."""

    #: Dimensionality of the vectors this provider emits.
    DIM: int

    @abstractmethod
    def embed(self, texts: Sequence[str]) -> np.ndarray:
        """Embed a batch of texts, returning an (N, DIM) float32 array."""
        raise NotImplementedError

    def embed_one(self, text: str) -> np.ndarray:
        """Embed a single text, returning a (DIM,) float32 array."""
        return self.embed([text])[0]


class HashingEmbedder(Embedder):
    """
    Deterministic, dependency-free embedding via the hashing trick.

    We tokenize into character 3-grams, hash each one into a bucket, and write
    +1/-1 (sign hashed from the gram) into a fixed-size vector, then L2-normalize.

    This is *not* semantically meaningful (synonyms are not mapped together), but
    it is stable, fast, and shares enough lexical overlap with the query for
    near-neighbors to cluster around the right topics.  Perfect for exercising
    the ANN pipeline without an API key.
    """

    DIM = HASHING_DIM
    _NGRAM = 3
    _SEED = 0xC0FFEE

    def __init__(self, dim: int = HASHING_DIM, ngram: int = 3):
        self.DIM = dim
        self._ngram = ngram

    @staticmethod
    def _grams(text: str, ngram: int = 3) -> list[str]:
        t = text.lower()
        return [t[i:i + ngram] for i in range(max(0, len(t) - ngram + 1))]

    def embed(self, texts: Sequence[str]) -> np.ndarray:
        out = np.zeros((len(texts), self.DIM), dtype=np.float32)
        ngram = self._ngram
        for i, text in enumerate(texts):
            for gram in self._grams(text or "", ngram):
                h = int.from_bytes(
                    hashlib.sha1(gram.encode("utf-8")).digest()[:8],
                    "little",
                )
                bucket = h % self.DIM
                sign = 1.0 if (h >> 63) & 1 else -1.0
                out[i, bucket] += sign
            norm = np.linalg.norm(out[i])
            if norm > 0:
                out[i] /= norm
        return out


class RemoteEmbedder(Embedder):
    """
    OpenAI-compatible remote embedding provider.

    Configured entirely through the environment:

        EMBEDDING_API_BASE  base URL, e.g. https://api.openai.com/v1
        EMBEDDING_API_KEY   API key (required)
        EMBEDDING_MODEL     model name, default text-embedding-3-small

    Any endpoint with the OpenAI embeddings shape works (OpenAI, Together,
    Groq, a local vLLM/Ollama server, etc.).
    """

    _DEFAULT_MODEL = "text-embedding-3-small"
    _DEFAULT_DIM = 1536  # text-embedding-3-small's default output dim

    def __init__(
        self,
        api_base: str | None = None,
        api_key: str | None = None,
        model: str | None = None,
        dim: int | None = None,
    ):
        import requests  # imported lazily so the fallback needs no network stack

        self._requests = requests
        self._api_base = (api_base or os.environ.get("EMBEDDING_API_BASE", "")).rstrip("/")
        self._api_key = api_key or os.environ.get("EMBEDDING_API_KEY", "")
        self._model = model or os.environ.get("EMBEDDING_MODEL", self._DEFAULT_MODEL)
        self._dim = dim or int(
            os.environ.get("EMBEDDING_DIM", self._DEFAULT_DIM)
        )
        if not self._api_base:
            raise ValueError(
                "EMBEDDING_API_BASE is not set; configure it (or use the local "
                "HashingEmbedder fallback)."
            )
        if not self._api_key:
            raise ValueError(
                "EMBEDDING_API_KEY is not set; set it, or use the local "
                "HashingEmbedder fallback."
            )

    @property
    def DIM(self) -> int:
        return self._dim

    def embed(self, texts: Sequence[str]) -> np.ndarray:
        resp = self._requests.post(
            f"{self._api_base}/embeddings",
            headers={"Authorization": f"Bearer {self._api_key}"},
            json={"model": self._model, "input": list(texts), "dimensions": self._dim},
            timeout=60,
        )
        resp.raise_for_status()
        payload = resp.json()
        items = payload["data"]
        # Some backends don't preserve input order; sort by the returned index.
        items = sorted(items, key=lambda d: d.get("index", 0))
        return np.array([d["embedding"] for d in items], dtype=np.float32)


def get_embedder(embedder: str | Embedder | None = None) -> Embedder:
    """
    Resolve an embedder from a name, an instance, or the environment.

    Resolution order:

      1. ``embedder="remote"`` or an ``Embedder`` instance -> use it directly.
      2. ``embedder="hashing"`` -> deterministic local fallback.
      3. If ``EMBEDDING_API_BASE`` and ``EMBEDDING_API_KEY`` are both set ->
         remote provider (the production path).
      4. Otherwise -> the local hashing fallback (with a printed note).

    This means the CLI works out of the box with no configuration: it silently
    degrades to the reproducible local embedding when no API key is present.
    """
    name = (embedder or "").lower() if isinstance(embedder, str) else None

    if isinstance(embedder, Embedder):
        return embedder
    if name == "hashing":
        return HashingEmbedder()
    if name == "remote":
        return RemoteEmbedder()

    api_base = os.environ.get("EMBEDDING_API_BASE")
    api_key = os.environ.get("EMBEDDING_API_KEY")
    if api_base and api_key:
        print(f"[embed] using remote provider: {api_base} (model={os.environ.get('EMBEDDING_MODEL')})")
        return RemoteEmbedder()

    print(
        "[embed] no EMBEDDING_API_BASE/KEY configured -> using local "
        "HashingEmbedder (deterministic, offline, non-semantic). Set the env vars "
        "or pass embedder=RemoteEmbedder() for real semantic embeddings."
    )
    return HashingEmbedder()