import contextlib
import os
import time

import httpx
from tenacity import RetryError, retry, retry_if_exception, stop_after_attempt, wait_exponential

EMBEDDING_MODEL = os.environ.get("OPENAI_EMBEDDING_MODEL", "gemini-embedding-2-preview")
BATCH_SIZE = int(os.environ.get("OPENAI_EMBEDDING_BATCH_SIZE", "2048"))
BATCH_DELAY = float(os.environ.get("OPENAI_EMBEDDING_BATCH_DELAY", "0"))


def _is_retryable(exc: BaseException) -> bool:
    if isinstance(exc, (httpx.ConnectError, httpx.TimeoutException)):
        return True
    if isinstance(exc, httpx.HTTPStatusError):
        status = exc.response.status_code
        return status == 429 or status >= 500
    return False


@retry(
    stop=stop_after_attempt(5),
    wait=wait_exponential(multiplier=1, min=2, max=30),
    retry=retry_if_exception(_is_retryable),
)
def _call_embedding_api(
    client: httpx.Client,
    url: str,
    headers: dict[str, str],
    payload: dict[str, str | list[str]],
) -> list[list[float]]:
    resp = client.post(url, json=payload, headers=headers)
    if resp.status_code == 429:
        retry_after = resp.headers.get("Retry-After")
        if retry_after is not None:
            with contextlib.suppress(ValueError):
                time.sleep(int(retry_after))
    resp.raise_for_status()
    data = resp.json()
    return [d["embedding"] for d in data["data"]]


def _embed_in_batches(
    client: httpx.Client,
    url: str,
    headers: dict[str, str],
    texts: list[str],
) -> list[list[float]]:
    n_batches = (len(texts) + BATCH_SIZE - 1) // BATCH_SIZE
    results: list[list[float]] = []
    for i in range(0, len(texts), BATCH_SIZE):
        batch = texts[i : i + BATCH_SIZE]
        payload: dict[str, str | list[str]] = {"model": EMBEDDING_MODEL, "input": batch}
        print(f"  batch {i // BATCH_SIZE + 1}/{n_batches}...", flush=True)
        results.extend(_call_embedding_api(client, url, headers, payload))
        if BATCH_DELAY > 0 and i + BATCH_SIZE < len(texts):
            time.sleep(BATCH_DELAY)
    return results


def embed_texts(texts: list[str]) -> list[list[float]] | None:
    if not texts:
        return []

    api_url = os.environ.get("OPENAI_BASE_URL")
    if not api_url:
        return None
    api_url = api_url.rstrip("/")
    api_key = os.environ.get("OPENAI_API_KEY", "")

    url = f"{api_url}/embeddings"
    headers = {"Authorization": f"Bearer {api_key}"}

    try:
        with httpx.Client(timeout=120.0) as client:
            if len(texts) <= BATCH_SIZE:
                payload: dict[str, str | list[str]] = {"model": EMBEDDING_MODEL, "input": texts}
                return _call_embedding_api(client, url, headers, payload)
            return _embed_in_batches(client, url, headers, texts)
    except (httpx.HTTPError, RetryError, ValueError, KeyError):
        return None
