"""Embedding helper for the ADE RAG playground.

Reuses quber's local `EmbeddingService` (BAAI/bge-large-en-v1.5, 1024 dims),
which runs on GPU when available and falls back to CPU. Both ingest (embed
each chunk) and query (embed the question) go through the same model, so the
vector spaces match.

Local embeddings are used instead of a hosted API so the experiment has no
external embedding dependency or quota.
"""

from __future__ import annotations

from functools import lru_cache
from typing import List

from quber.db.embeddings import get_embedding_service

EMBED_DIM = 1024  # bge-large-en-v1.5


@lru_cache(maxsize=1)
def _service():
    # The provider is always local, so EMBEDDING_PROVIDER is ignored here. The
    # device comes from EMBEDDING_DEVICE, default cuda, and a cuda model falls
    # back to CPU when CUDA is not available.
    return get_embedding_service(provider="local")


def embed_texts(texts: List[str]) -> List[List[float]]:
    if not texts:
        return []
    arr = _service().embed(texts, batch_size=32)
    return [row.tolist() for row in arr]


def embed_query(text: str) -> List[float]:
    return embed_texts([text])[0]


def embed_document(doc_key: str, texts: List[str]) -> List[List[float]]:
    """Embed a document's chunks: on the GPU worker when this host is set up
    to send work there, locally otherwise. The question at query time always
    embeds locally; only ingestion, thousands of chunks at once, is worth a
    worker."""
    from quber.playground import gpu

    if gpu.configured():
        return gpu.embed(doc_key, texts)
    return embed_texts(texts)