"""
SiliconFlow Embedding 实现
参考 aigroup-llamaindex 项目的实现
"""
from typing import List, Any, Optional
from loguru import logger

import httpx
from llama_index.core.embeddings import BaseEmbedding


class SiliconFlowEmbedding(BaseEmbedding):
    """SiliconFlow 嵌入模型实现"""
    
    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.siliconflow.cn/v1",
        model_name: str = "Qwen/Qwen3-Embedding-8B",
        timeout: float = 60.0,
        **kwargs: Any
    ):
        """初始化 SiliconFlow 嵌入模型
        
        Args:
            api_key: API 密钥
            base_url: API 基础 URL
            model_name: 嵌入模型名称
            timeout: 请求超时时间（秒）
        """
        super().__init__(model_name=model_name, **kwargs)
        self._api_key = api_key
        self._base_url = base_url
        self._model_name = model_name
        self._timeout = timeout
        
        # 初始化 HTTP 客户端
        self._client = httpx.Client(
            timeout=self._timeout,
            limits=httpx.Limits(max_connections=10, max_keepalive_connections=5)
        )
        
        masked_key = f"{api_key[:6]}...{api_key[-4:]}" if len(api_key) > 10 else "****"
        logger.info(f"初始化 SiliconFlow 嵌入模型: {model_name}, API Key: {masked_key}")
    
    @classmethod
    def class_name(cls) -> str:
        """返回类名"""
        return "SiliconFlowEmbedding"
    
    def _get_headers(self) -> dict:
        """获取 HTTP 请求头"""
        return {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json"
        }
    
    def _get_embedding(self, text: str) -> List[float]:
        """获取单个文本的嵌入向量"""
        return self._get_embeddings([text])[0]
    
    def _get_embeddings(self, texts: List[str]) -> List[List[float]]:
        """批量获取嵌入向量（同步）"""
        if not texts:
            return []
        
        url = f"{self._base_url}/embeddings"
        data = {
            "input": texts,
            "model": self._model_name
        }
        
        try:
            response = self._client.post(url, json=data, headers=self._get_headers())
            response.raise_for_status()
            
            result = response.json()
            embeddings = [item["embedding"] for item in result["data"]]
            
            return embeddings
        except Exception as e:
            logger.error(f"获取嵌入向量失败: {e}")
            # 如果批量请求失败，回退到逐个请求
            logger.warning("批量请求失败，回退到逐个请求")
            embeddings = []
            for text in texts:
                try:
                    single_data = {
                        "input": text,
                        "model": self._model_name
                    }
                    response = self._client.post(url, json=single_data, headers=self._get_headers())
                    response.raise_for_status()
                    result = response.json()
                    embeddings.append(result["data"][0]["embedding"])
                except Exception as single_e:
                    logger.error(f"获取单个文本嵌入向量失败: {single_e}")
                    # 返回零向量作为占位符
                    embeddings.append([0.0] * 1024)  # 假设维度为1024
            return embeddings
    
    async def _aget_query_embedding(self, query: str) -> List[float]:
        """异步获取查询文本的嵌入向量"""
        return self._get_embedding(query)
    
    async def _aget_text_embedding(self, text: str) -> List[float]:
        """异步获取文本的嵌入向量"""
        return self._get_embedding(text)
    
    def _get_query_embedding(self, query: str) -> List[float]:
        """同步获取查询文本的嵌入向量"""
        return self._get_embedding(query)
    
    def _get_text_embedding(self, text: str) -> List[float]:
        """同步获取文本的嵌入向量"""
        return self._get_embedding(text)
    
    async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        """批量获取嵌入向量（异步）"""
        # 异步版本暂时使用同步方法
        return self._get_embeddings(texts)
    
    def get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        """批量获取嵌入向量（同步）"""
        return self._get_embeddings(texts)
    
    def close(self) -> None:
        """关闭 HTTP 客户端"""
        if self._client:
            self._client.close()
            self._client = None
            logger.info("SiliconFlow HTTP 客户端已关闭")
    
    def __del__(self):
        """清理资源"""
        try:
            if self._client:
                self._client.close()
        except Exception:
            pass
