from __future__ import annotations
import hashlib
import threading
import time
import httpx

VOYAGE_API_URL = "https://api.voyageai.com/v1/embeddings"

# LRU-style cache for query embeddings. Keyed by (text, model) hash.
# Avoids redundant Voyage AI API calls when multiple concurrent agents
# query the same drug/condition terms.
_embedding_cache: dict[str, list[float]] = {}
_cache_lock = threading.Lock()
_CACHE_MAX_SIZE = 512


def embed_texts(
    texts: list[str],
    api_key: str,
    model: str = "voyage-3",
    input_type: str = "document",
    batch_size: int = 128,
    max_retries: int = 3,
    timeout: float = 60.0,
) -> list[list[float]]:
    """Embed a list of texts using the Voyage AI API."""
    all_embeddings: list[list[float]] = []

    for i in range(0, len(texts), batch_size):
        # Proactive rate limiting: 3 req/min = 1 req per 20s
        if i > 0:
            time.sleep(21)
        batch = texts[i : i + batch_size]
        for attempt in range(max_retries):
            response = httpx.post(
                VOYAGE_API_URL,
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                json={
                    "input": batch,
                    "model": model,
                    "input_type": input_type,
                },
                timeout=timeout,
            )
            if response.status_code == 429:
                wait = 30 * (attempt + 1)
                time.sleep(wait)
                continue
            response.raise_for_status()
            break
        else:
            response.raise_for_status()
        data = response.json()["data"]
        all_embeddings.extend(item["embedding"] for item in data)

    return all_embeddings


def embed_query(
    text: str,
    api_key: str,
    model: str = "voyage-3",
) -> list[float]:
    """Embed a single query text for similarity search.

    Uses a shorter timeout (10s) and no retries since this is called
    during interactive queries where latency matters.

    Results are cached in-memory to avoid redundant API calls when
    multiple concurrent agents query the same terms.
    """
    cache_key = hashlib.sha256(f"{model}:{text}".encode()).hexdigest()

    with _cache_lock:
        if cache_key in _embedding_cache:
            return _embedding_cache[cache_key]

    results = embed_texts(
        [text], api_key=api_key, model=model,
        input_type="query", max_retries=1, timeout=10.0,
    )
    embedding = results[0]

    with _cache_lock:
        if len(_embedding_cache) >= _CACHE_MAX_SIZE:
            # Evict oldest entry (first inserted)
            oldest = next(iter(_embedding_cache))
            del _embedding_cache[oldest]
        _embedding_cache[cache_key] = embedding

    return embedding
