import time
from typing import List

import httpx


class OllamaEmbedder:
    def __init__(
        self,
        base_url: str,
        model: str,
        batch_size: int,
        timeout: float = 300.0,
    ):
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.batch_size = batch_size
        self.timeout = timeout
        self._client = httpx.Client(timeout=timeout)

    def embed(
        self, texts: List[str], retries: int = 3, backoff: float = 1.0
    ) -> List[List[float]]:
        all_embeddings: List[List[float]] = []
        url = f"{self.base_url}/api/embed"

        for i in range(0, len(texts), self.batch_size):
            batch = texts[i : i + self.batch_size]
            last_exception: Exception | None = None

            for attempt in range(retries):
                try:
                    response = self._client.post(
                        url,
                        json={"model": self.model, "input": batch},
                    )
                    response.raise_for_status()
                    data = response.json()
                    embeddings = data["embeddings"]

                    # Validate response count matches batch count
                    if len(embeddings) != len(batch):
                        raise ValueError(
                            f"Ollama returned {len(embeddings)} embeddings for {len(batch)} texts"
                        )

                    all_embeddings.extend(embeddings)
                    break
                except (httpx.HTTPError, httpx.RequestError) as exc:
                    last_exception = exc
                    if attempt < retries - 1:
                        time.sleep(backoff * (2**attempt))
            else:
                raise RuntimeError(
                    f"Failed to embed batch after {retries} retries. "
                    f"Last error: {last_exception}"
                ) from last_exception

        return all_embeddings

    def close(self) -> None:
        """Close the HTTP client."""
        self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False
