"""OpenAI embedding wrapper.

支持 OpenAI 官方 API、兼容 API 以及 sagellm 服务端 embedding。
"""

import logging
import os
from typing import TYPE_CHECKING, Any

from ..base import BaseEmbedding

if TYPE_CHECKING:
    pass  # Reserved for future type hints

# 抑制 httpx 的 INFO 日志（每次 HTTP 请求都会打印）
logging.getLogger("httpx").setLevel(logging.WARNING)

logger = logging.getLogger(__name__)


class OpenAIEmbedding(BaseEmbedding):
    """OpenAI Embedding API Wrapper

    支持 OpenAI 官方 API、兼容的第三方 API（如 DeepSeek 等）
    以及 sagellm 服务端 embedding。

    特点:
        - ✅ 高质量 embedding
        - ✅ 支持多种模型
        - ✅ 兼容第三方 API
        - ✅ 支持 sagellm 服务端 embedding
        - ❌ 需要 API Key（仅 openai provider）
        - ❌ 需要网络连接（仅 openai provider）
        - 💰 按使用量计费（仅 openai provider）

    支持的模型:
        - text-embedding-3-small (1536维，性价比高)
        - text-embedding-3-large (3072维，最高质量)
        - text-embedding-ada-002 (1536维，旧版本)
        - 任意本地/远程模型（通过 sagellm provider）

    Args:
        model: 模型名称（默认 'text-embedding-3-small'）
        api_key: API 密钥（可选，默认从环境变量 OPENAI_API_KEY 读取）
        base_url: API 端点（可选，用于兼容 API）
        provider: 后端提供者，'openai'（默认）或 'sagellm'（服务端）
        sagellm_config: sagellm 服务配置（仅 provider='sagellm' 时有效）

    Examples:
        >>> # OpenAI 官方 API
        >>> import os
        >>> emb = OpenAIEmbedding(
        ...     model="text-embedding-3-small",
        ...     api_key=os.getenv("OPENAI_API_KEY")
        ... )
        >>> vec = emb.embed("hello world")
        >>>
        >>> # 兼容 API（自定义端点）
        >>> emb = OpenAIEmbedding(
        ...     model="text-embedding-v1",
        ...     api_key=os.getenv("OPENAI_API_KEY"),
        ...     base_url=os.getenv("OPENAI_BASE_URL", "http://localhost:8090/v1")
        ... )
        >>> vec = emb.embed("你好世界")
        >>>
        >>> # OpenAI 兼容服务部署的模型
        >>> emb = OpenAIEmbedding(
        ...     model="BAAI/bge-base-en-v1.5",
        ...     base_url="http://localhost:8000/v1"
        ... )
        >>>
        >>> # sagellm 服务端 embedding
        >>> emb = OpenAIEmbedding(
        ...     model="BAAI/bge-small-zh-v1.5",
        ...     provider="sagellm",
        ...     sagellm_config={"base_url": "http://localhost:8000"}
        ... )
        >>> vec = emb.embed("你好世界")
    """

    # 常见模型的维度映射
    DIMENSION_MAP = {
        "text-embedding-3-small": 1536,
        "text-embedding-3-large": 3072,
        "text-embedding-ada-002": 1536,
        "text-embedding-v1": 1536,
        # Common model aliases for sagellm
        "BAAI/bge-small-zh-v1.5": 512,
        "BAAI/bge-base-zh-v1.5": 768,
        "BAAI/bge-large-zh-v1.5": 1024,
        "BAAI/bge-small-en-v1.5": 384,
        "BAAI/bge-base-en-v1.5": 768,
        "BAAI/bge-large-en-v1.5": 1024,
        "BAAI/bge-m3": 1024,
    }

    # 支持的 provider
    SUPPORTED_PROVIDERS = ("openai", "sagellm")

    def __init__(
        self,
        model: str = "text-embedding-3-small",
        api_key: str | None = None,
        base_url: str | None = None,
        provider: str = "openai",
        sagellm_config: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        """初始化 OpenAI Embedding

        Args:
            model: 模型名称
            api_key: API 密钥（可选，仅 openai provider）
            base_url: API 端点（可选，仅 openai provider）
            provider: 后端提供者，'openai' 或 'sagellm'
            sagellm_config: sagellm 推理配置（仅 provider='sagellm' 时有效）
            **kwargs: 其他参数（保留用于扩展）

        Raises:
            RuntimeError: 如果 openai provider 未提供 API Key
            ValueError: 如果 provider 不支持
        """
        super().__init__(model=model, api_key=api_key, base_url=base_url, **kwargs)

        self._model = model
        self._provider = provider.lower()
        self._sagellm_config = sagellm_config or {}
        self._sagellm_client_cfg: dict[str, Any] | None = None

        # 验证 provider
        if self._provider not in self.SUPPORTED_PROVIDERS:
            raise ValueError(
                f"不支持的 provider: {self._provider}\n"
                f"支持的 provider: {', '.join(self.SUPPORTED_PROVIDERS)}"
            )

        if self._provider == "openai":
            # OpenAI API 模式
            self._api_key = api_key or os.getenv("OPENAI_API_KEY")
            self._base_url = base_url

            # 检查 API Key
            if not self._api_key:
                raise RuntimeError(
                    "OpenAI embedding 需要 API Key。\n"
                    "解决方案:\n"
                    "  1. 设置环境变量: export OPENAI_API_KEY='your-key'\n"  # pragma: allowlist secret
                    "  2. 传递参数: OpenAIEmbedding(api_key='your-key', ...)\n"  # pragma: allowlist secret
                    "\n"
                    "如果使用兼容 API:\n"
                    "  export OPENAI_API_KEY='your-api-key'\n"  # pragma: allowlist secret
                    "  并指定 base_url 参数\n"
                    "\n"
                    "或者使用 sagellm 服务:\n"
                    "  OpenAIEmbedding(model='BAAI/bge-small-zh-v1.5', provider='sagellm')"
                )
        else:
            # sagellm 服务模式
            self._api_key = None
            self._base_url = None
            self._sagellm_client_cfg = self._build_sagellm_client_cfg()
            logger.info(
                f"使用 sagellm 服务端 embedding: model={model}, config={self._sagellm_config}"
            )

        # 推断或获取维度
        self._dim = self._infer_dimension()

    def embed(self, text: str) -> list[float]:
        """将文本转换为 embedding 向量

        Args:
            text: 输入文本

        Returns:
            embedding 向量

        Raises:
            RuntimeError: 如果 API 调用或本地推理失败
        """
        if self._provider == "sagellm":
            return self._embed_with_sagellm(text)
        return self._embed_with_openai(text)

    def _embed_with_openai(self, text: str) -> list[float]:
        """使用 OpenAI API 生成 embedding

        Args:
            text: 输入文本

        Returns:
            embedding 向量
        """
        try:
            from openai import OpenAI

            client = OpenAI(api_key=self._api_key, base_url=self._base_url)
            response = client.embeddings.create(
                model=self._model,
                input=text,
            )
            return response.data[0].embedding
        except Exception as e:
            raise RuntimeError(
                f"OpenAI embedding 失败: {e}\n"
                f"模型: {self._model}\n"
                f"文本: {text[:100]}...\n"
                f"提示: 检查 API Key 是否有效，网络连接是否正常"
            ) from e

    def _embed_with_sagellm(self, text: str) -> list[float]:
        """使用 sagellm 服务生成 embedding

        Args:
            text: 输入文本

        Returns:
            embedding 向量
        """
        try:
            from sage.common.components.sage_embedding import sagellm

            client_cfg = self._sagellm_client_cfg or self._build_sagellm_client_cfg()
            return sagellm.sagellm_embed_sync(
                text,
                tokenizer=self._model,
                embed_model=client_cfg,
            )
        except Exception as e:
            raise RuntimeError(
                f"sagellm embedding 失败: {e}\n"
                f"模型: {self._model}\n"
                f"文本: {text[:100]}...\n"
                "提示: 检查 sagellm 服务是否可用、SAGELLM_BASE_URL 是否正确"
            ) from e

    def _build_sagellm_client_cfg(self) -> dict[str, Any]:
        """构建 sagellm embedding 客户端配置。"""
        from sage.common.components.sage_embedding import sagellm

        cfg = sagellm.initialize_sagellm_client(self._model)
        if "base_url" in self._sagellm_config:
            cfg["base_url"] = str(self._sagellm_config["base_url"]).rstrip("/")
        if "api_key" in self._sagellm_config:
            cfg["api_key"] = self._sagellm_config["api_key"]
        if "timeout_s" in self._sagellm_config:
            cfg["timeout_s"] = float(self._sagellm_config["timeout_s"])
        return cfg

    def embed_batch(self, texts: list[str]) -> list[list[float]]:
        """批量将文本转换为 embedding 向量

        使用 OpenAI API 的批量接口或 sagellm 的批量编码。

        Args:
            texts: 输入文本列表

        Returns:
            embedding 向量列表

        Raises:
            RuntimeError: 如果 API 调用或本地推理失败
        """
        if not texts:
            return []

        if self._provider == "sagellm":
            return self._embed_batch_with_sagellm(texts)
        return self._embed_batch_with_openai(texts)

    def _embed_batch_with_openai(self, texts: list[str]) -> list[list[float]]:
        """使用 OpenAI API 批量生成 embedding

        Args:
            texts: 输入文本列表

        Returns:
            embedding 向量列表
        """
        try:
            from openai import OpenAI

            # 设置环境变量
            if self._api_key:
                import os

                os.environ["OPENAI_API_KEY"] = self._api_key

            client = OpenAI(base_url=self._base_url)

            # OpenAI API 支持批量：input 可以是字符串列表
            response = client.embeddings.create(
                model=self._model,
                input=texts,  # 直接传入列表
            )

            # 按照原始顺序返回结果
            return [item.embedding for item in response.data]

        except Exception as e:
            raise RuntimeError(
                f"OpenAI 批量 embedding 失败: {e}\n"
                f"模型: {self._model}\n"
                f"批量大小: {len(texts)}\n"
                f"提示: 检查 API Key 是否有效，网络连接是否正常"
            ) from e

    def _embed_batch_with_sagellm(self, texts: list[str]) -> list[list[float]]:
        """使用 sagellm 服务批量生成 embedding

        Args:
            texts: 输入文本列表

        Returns:
            embedding 向量列表
        """
        try:
            from sage.common.components.sage_embedding import sagellm

            client_cfg = self._sagellm_client_cfg or self._build_sagellm_client_cfg()
            return sagellm.sagellm_embed_batch_sync(
                texts,
                tokenizer=self._model,
                embed_model=client_cfg,
            )
        except Exception as e:
            raise RuntimeError(
                f"sagellm 批量 embedding 失败: {e}\n"
                f"模型: {self._model}\n"
                f"批量大小: {len(texts)}\n"
                "提示: 检查 sagellm 服务是否可用、SAGELLM_BASE_URL 是否正确"
            ) from e

    def get_dim(self) -> int:
        """获取向量维度

        Returns:
            维度值
        """
        return self._dim

    @property
    def method_name(self) -> str:
        """返回方法名称

        Returns:
            'openai' 或 'sagellm'（取决于 provider）
        """
        return self._provider

    @property
    def provider(self) -> str:
        """返回后端提供者

        Returns:
            'openai' 或 'sagellm'
        """
        return self._provider

    def _infer_dimension(self) -> int:
        """推断或获取维度

        Returns:
            推断的维度值
        """
        # 优先使用已知的维度映射
        if self._model in self.DIMENSION_MAP:
            return self.DIMENSION_MAP[self._model]

        # 如果是未知模型，尝试通过实际调用推断
        try:
            sample = self.embed("test")
            return len(sample)
        except Exception:
            # 如果推断失败，返回默认维度
            return 1536

    @classmethod
    def get_model_info(cls) -> dict[str, Any]:
        """返回模型元信息

        Returns:
            模型信息字典
        """
        return {
            "method": "openai",
            "requires_api_key": True,  # Only for openai provider
            "requires_model_download": False,
            "default_dimension": 1536,
            "supported_providers": list(cls.SUPPORTED_PROVIDERS),
        }

    def __repr__(self) -> str:
        """返回对象的字符串表示

        Returns:
            字符串表示
        """
        base_info = (
            f"OpenAIEmbedding(model='{self._model}', dim={self._dim}, provider='{self._provider}'"
        )
        if self._provider == "openai" and self._base_url:
            base_info += f", base_url='{self._base_url}'"
        elif self._provider == "sagellm" and self._sagellm_config:
            config_str = ", ".join(f"{k}={v!r}" for k, v in self._sagellm_config.items())
            base_info += f", config={{{config_str}}}"
        return base_info + ")"

    def close(self) -> None:
        """释放资源

        当前为服务端调用，无本地模型资源需要释放。
        """
        self._sagellm_client_cfg = None

    def __del__(self) -> None:
        """析构函数"""
        try:
            self.close()
        except Exception:
            pass  # 忽略析构时的错误
