"""sagellm service-backed embedding helpers."""

from __future__ import annotations

import json
import os
from functools import lru_cache
from typing import Any
from urllib import error as urllib_error
from urllib import request as urllib_request


def _normalize_base_url(base_url: str) -> str:
    normalized = base_url.rstrip("/")
    if normalized.endswith("/v1"):
        return normalized[:-3]
    return normalized


@lru_cache(maxsize=8)
def initialize_sagellm_client(model_name: str) -> dict[str, Any]:
    """Initialize sagellm embedding client configuration."""
    base_url = _normalize_base_url(os.getenv("SAGELLM_BASE_URL", "http://localhost:8000"))
    api_key = os.getenv("SAGELLM_API_KEY") or os.getenv("OPENAI_API_KEY") or ""
    return {
        "kind": "sagellm_embedding_client",
        "model": model_name,
        "base_url": base_url,
        "api_key": api_key,
        "timeout_s": float(os.getenv("SAGELLM_TIMEOUT_S", "60")),
    }


def _post_embeddings(client_cfg: dict[str, Any], texts: list[str]) -> list[list[float]]:
    payload: dict[str, Any] = {
        "model": client_cfg["model"],
        "input": texts,
    }
    headers = {"Content-Type": "application/json"}
    if client_cfg.get("api_key"):
        headers["Authorization"] = f"Bearer {client_cfg['api_key']}"

    url = f"{client_cfg['base_url']}/v1/embeddings"
    req = urllib_request.Request(
        url=url,
        data=json.dumps(payload).encode("utf-8"),
        headers=headers,
        method="POST",
    )

    try:
        with urllib_request.urlopen(req, timeout=float(client_cfg.get("timeout_s", 60.0))) as resp:  # noqa: S310
            data = json.loads(resp.read().decode("utf-8"))
    except urllib_error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"sagellm embedding HTTP {exc.code}: {body}") from exc
    except urllib_error.URLError as exc:
        raise RuntimeError(
            f"Cannot connect to sagellm at {client_cfg['base_url']}: {exc.reason}"
        ) from exc

    items = data.get("data", [])
    if not items:
        return []
    return [item.get("embedding", []) for item in items]


def _resolve_client_cfg(embed_model: Any, model_name: str | None = None) -> dict[str, Any]:
    if isinstance(embed_model, dict) and embed_model.get("kind") == "sagellm_embedding_client":
        return embed_model
    if model_name:
        return initialize_sagellm_client(model_name)
    raise RuntimeError(
        "sagellm embedding requires backend configuration. "
        "Please provide model or embed_model client config."
    )


def sagellm_embed_sync(text: str, tokenizer: Any = None, embed_model: Any = None) -> list[float]:
    """Generate single embedding through sagellm service."""
    if not text:
        return []
    model_name = tokenizer if isinstance(tokenizer, str) else None
    client_cfg = _resolve_client_cfg(embed_model, model_name=model_name)
    embeddings = _post_embeddings(client_cfg, [text])
    return embeddings[0] if embeddings else []


def sagellm_embed_batch_sync(
    texts: list[str], tokenizer: Any = None, embed_model: Any = None
) -> list[list[float]]:
    """Generate batch embeddings through sagellm service."""
    if not texts:
        return []
    model_name = tokenizer if isinstance(tokenizer, str) else None
    client_cfg = _resolve_client_cfg(embed_model, model_name=model_name)
    return _post_embeddings(client_cfg, texts)
