from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Awaitable, Callable

import httpx

EmbedFn = Callable[[list[str]], Awaitable[list[list[float]]]]


@dataclass
class Embedder:
    fn: EmbedFn
    dimension: int


def make_openai_compatible_embed_fn(api_key: str, model: str, base_url: str, dimension: int) -> Embedder:
    url = base_url.rstrip("/") + "/embeddings"

    async def embed(texts: list[str]) -> list[list[float]]:
        async with httpx.AsyncClient() as client:
            response = await client.post(
                url,
                headers={"Authorization": f"Bearer {api_key}"},
                json={"input": texts, "model": model},
            )
            response.raise_for_status()
            data = response.json()
            items = sorted(data["data"], key=lambda x: x["index"])
            return [item["embedding"] for item in items]

    return Embedder(fn=embed, dimension=dimension)


def get_embedder(provider: str, model: str) -> Embedder:
    if provider in ("default", "openai"):
        return make_openai_compatible_embed_fn(
            api_key=os.environ["DOT__EMBED_API_KEY"],
            model=os.environ["DOT__EMBED_MODEL"] if model in ("default", "") else model,
            base_url=os.environ["DOT__EMBED_BASE_URL"],
            dimension=int(os.environ["DOT__EMBED_DIMENSION"]),
        )

    raise ValueError(f"Unknown provider: {provider!r}. Supported: 'default', 'openai'")
