"""Supply zembed's saved prompt and terminal suffix to Infinity's /embed API."""
import json
import urllib.error
import urllib.request


def embed(texts, kind="document", base_url="http://localhost:7997", model="michaelfeil/zembed-1-infinity"):
    if isinstance(texts, str):
        texts = [texts]
    texts = list(texts)
    if kind not in ("query", "document") or not texts or any(not isinstance(text, str) for text in texts):
        raise ValueError("Provide text strings and kind='query' or kind='document'")
    body = {
        "model": model,
        "inputs": [text + "<|im_end|>\n" for text in texts],
        "prompt_name": kind,
        "normalize": True,
        "truncate": False,
    }
    request = urllib.request.Request(
        base_url.rstrip("/") + "/embed",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(request, timeout=300) as response:
            return json.load(response)
    except urllib.error.HTTPError as error:
        raise RuntimeError(f"HTTP {error.code}: {error.read().decode()}") from error