"""sagellm service-backed embedding wrapper."""

from __future__ import annotations

from typing import Any

from ..base import BaseEmbedding
from ..sagellm import (
    initialize_sagellm_client,
    sagellm_embed_batch_sync,
    sagellm_embed_sync,
)


class SagellmEmbedding(BaseEmbedding):
    """Embedding wrapper backed by sagellm OpenAI-compatible endpoint."""

    def __init__(self, model: str, **kwargs: Any) -> None:
        super().__init__(model=model, **kwargs)

        try:
            self.tokenizer = model
            self.embed_model = initialize_sagellm_client(model)
        except Exception as error:
            raise RuntimeError(
                f"Failed to initialize sagellm embedding backend for '{model}': {error}\n"
                "提示:\n"
                "  1. 检查 SAGELLM_BASE_URL 是否正确\n"
                "  2. 确认 sagellm 服务已启动\n"
                "  3. 检查模型在 sagellm 侧是否可用"
            ) from error

        self._dim = self._infer_dimension()
        self._model_name = model

    def embed(self, text: str) -> list[float]:
        try:
            return sagellm_embed_sync(text, self.tokenizer, self.embed_model)
        except Exception as error:
            raise RuntimeError(f"sagellm embedding 失败: {error}\n文本: {text[:100]}...") from error

    def embed_batch(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            return []

        try:
            return sagellm_embed_batch_sync(texts, self.tokenizer, self.embed_model)
        except Exception as error:
            raise RuntimeError(
                f"sagellm batch embedding 失败: {error}\n文本数量: {len(texts)}"
            ) from error

    def get_dim(self) -> int:
        return self._dim

    @property
    def method_name(self) -> str:
        return "sagellm"

    def _infer_dimension(self) -> int:
        try:
            sample = self.embed("test")
            return len(sample)
        except Exception:
            return 768

    @classmethod
    def get_model_info(cls) -> dict[str, Any]:
        return {
            "method": "sagellm",
            "requires_api_key": False,
            "requires_model_download": False,
            "default_dimension": None,
        }

    def __repr__(self) -> str:
        return f"SagellmEmbedding(model='{self._model_name}', dim={self._dim})"
