"""
Embedding 服务

通过 OpenAI 兼容的 Embeddings API 将文本转换为向量
支持所有兼容 OpenAI 接口的提供商（OpenAI、Qwen、Ollama 等）
"""
import logging
import math
from typing import List, Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)

# 单次批量请求的默认最大文本数
DEFAULT_BATCH_SIZE = 50

# 不同提供商的批次大小限制
PROVIDER_BATCH_SIZE = {
    'qwen': 10,           # 阿里云 DashScope 限制 10
    'dashscope': 10,
    'siliconflow': 10,    # 硅基流动限制较小
    'ollama': 1,          # Ollama 通常逐条处理
}



class EmbeddingService:
    """
    Embedding 服务
    
    通过模型 ID 获取对应的提供商，调用 OpenAI 兼容的 Embeddings API
    """

    def __init__(self, db: AsyncSession):
        self._db = db
        self._client_cache = {}

    async def _get_client_and_model(self, model_id: str):
        """
        根据模型 ID 获取异步客户端和模型名称
        
        Returns:
            (async_client, model_name, max_tokens, provider_type)
        """
        from ai_platform.models import LLMModel, LLMProvider

        result = await self._db.execute(
            select(LLMModel).where(
                LLMModel.id == model_id,
                LLMModel.is_active == True,
                LLMModel.is_deleted == False
            )
        )
        model = result.scalar_one_or_none()
        if not model:
            raise ValueError(f'Embedding 模型不存在或已禁用: {model_id}')
        if model.model_type != 'embedding':
            raise ValueError(f'模型 {model.display_name} 不是 Embedding 类型')

        provider_result = await self._db.execute(
            select(LLMProvider).where(
                LLMProvider.id == model.provider_id,
                LLMProvider.is_active == True,
                LLMProvider.is_deleted == False
            )
        )
        provider = provider_result.scalar_one_or_none()
        if not provider:
            raise ValueError('Embedding 模型对应的提供商不存在或已禁用')

        cache_key = str(provider.id)
        if cache_key not in self._client_cache:
            import httpx
            from openai import AsyncOpenAI

            # 根据提供商类型确定 base_url
            if provider.provider_type == 'ollama':
                base_url = (provider.ollama_host or 'http://localhost:11434').rstrip('/') + '/v1'
            else:
                from ai_platform.providers.registry import ProviderRegistry
                provider_class = ProviderRegistry.get(provider.provider_type)
                default_base = getattr(provider_class, 'DEFAULT_API_BASE', 'https://api.openai.com/v1') if provider_class else 'https://api.openai.com/v1'
                base_url = provider.api_base or default_base

            self._client_cache[cache_key] = AsyncOpenAI(
                api_key=provider.api_key or 'ollama',
                base_url=base_url,
                timeout=httpx.Timeout(120.0, connect=30.0),
                max_retries=5,
            )

        return self._client_cache[cache_key], model.model_name, model.context_window or 8191, provider.provider_type

    async def embed_text(
        self,
        model_id: str,
        text: str,
        dimensions: Optional[int] = None,
    ) -> List[float]:
        """
        将单个文本转换为向量
        
        Args:
            model_id: Embedding 模型 ID
            text: 文本内容
            dimensions: 向量维度（可选，部分模型支持）
            
        Returns:
            向量列表 List[float]
        """
        results = await self.embed_texts(model_id, [text], dimensions)
        return results[0]

    async def embed_texts(
        self,
        model_id: str,
        texts: List[str],
        dimensions: Optional[int] = None,
    ) -> List[List[float]]:
        """
        批量将文本转换为向量
        
        Args:
            model_id: Embedding 模型 ID
            texts: 文本列表
            dimensions: 向量维度（可选）
            
        Returns:
            向量列表 List[List[float]]
        """
        if not texts:
            return []

        client, model_name, max_tokens, provider_type = await self._get_client_and_model(model_id)

        # 根据提供商类型确定批次大小
        batch_size = PROVIDER_BATCH_SIZE.get(provider_type, DEFAULT_BATCH_SIZE)

        # 预处理：截断超长文本
        processed_texts = []
        for text in texts:
            if not text or not text.strip():
                processed_texts.append(" ")
            else:
                # 粗略估算 token 数（中文约 1 字 = 1.5 token，英文约 4 字符 = 1 token）
                # 保守截断到 max_tokens * 2 个字符
                max_chars = max_tokens * 2
                if len(text) > max_chars:
                    processed_texts.append(text[:max_chars])
                else:
                    processed_texts.append(text)

        # 分批处理
        all_embeddings = [None] * len(processed_texts)
        total_batches = math.ceil(len(processed_texts) / batch_size)

        for batch_idx in range(total_batches):
            start = batch_idx * batch_size
            end = min(start + batch_size, len(processed_texts))
            batch_texts = processed_texts[start:end]

            try:
                kwargs = {
                    'model': model_name,
                    'input': batch_texts,
                }
                # 仅对明确支持 dimensions 参数的模型传递该参数
                if dimensions:
                    model_lower = model_name.lower()
                    # OpenAI text-embedding-3 系列原生支持任意 dimensions
                    if 'text-embedding-3' in model_lower:
                        kwargs['dimensions'] = dimensions
                    # DashScope text-embedding-v3 只接受 [64,128,256,512,768,1024]
                    elif 'text-embedding-v3' in model_lower and dimensions in (64, 128, 256, 512, 768, 1024):
                        kwargs['dimensions'] = dimensions

                response = await client.embeddings.create(**kwargs)

                for item in response.data:
                    all_embeddings[start + item.index] = item.embedding

            except Exception as e:
                logger.exception(f'Embedding 批次 {batch_idx + 1}/{total_batches} 失败: {e}')
                raise ValueError(f'Embedding 调用失败: {str(e)}')

        # 检查是否所有向量都已生成
        for i, emb in enumerate(all_embeddings):
            if emb is None:
                raise ValueError(f'第 {i} 个文本的向量未生成')

        return all_embeddings

    async def get_embedding_dimensions(self, model_id: str) -> int:
        """
        获取模型的向量维度（通过嵌入一个测试文本来检测）
        
        Args:
            model_id: Embedding 模型 ID
            
        Returns:
            向量维度
        """
        test_embedding = await self.embed_text(model_id, "test")
        return len(test_embedding)
