"""
嵌入服务模块
支持多种嵌入模型提供商：OpenAI / 智谱 / 阿里云 DashScope / 本地 BGE
"""

from __future__ import annotations

import abc
import logging
from typing import List, Optional

import numpy as np

from app.config import EmbeddingProvider, settings

logger = logging.getLogger(__name__)


# ──────────────────────────── 抽象基类 ────────────────────────────

class EmbeddingProviderBase(abc.ABC):
    """嵌入模型提供商抽象基类"""

    @abc.abstractmethod
    async def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """
        批量生成文本嵌入向量。

        Args:
            texts: 待嵌入的文本列表

        Returns:
            嵌入向量列表，每个向量是 float 的列表

        Raises:
            Exception: 嵌入生成失败时抛出异常
        """
        ...

    @abc.abstractmethod
    async def embed_single(self, text: str) -> List[float]:
        """生成单条文本的嵌入向量"""
        ...

    @property
    @abc.abstractmethod
    def dimension(self) -> int:
        """嵌入向量维度"""
        ...


# ──────────────────────────── OpenAI 实现 ────────────────────────────

class OpenAIEmbedding(EmbeddingProviderBase):
    """OpenAI 嵌入模型（兼容 OpenAI API 格式的服务）"""

    def __init__(self) -> None:
        from openai import AsyncOpenAI

        self._client = AsyncOpenAI(
            api_key=settings.OPENAI_API_KEY or "EMPTY",
            base_url=settings.OPENAI_BASE_URL,
        )
        self._model = settings.EMBEDDING_MODEL or "text-embedding-3-small"

    @property
    def dimension(self) -> int:
        return settings.EMBEDDING_DIMENSIONS

    async def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """调用 OpenAI 兼容接口批量生成嵌入"""
        # OpenAI 接口单次最多 2048 条，自动分批
        batch_size = 2048
        all_embeddings: List[List[float]] = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            response = await self._client.embeddings.create(
                input=batch,
                model=self._model,
                dimensions=self.dimension if self.dimension <= 3072 else None,
            )
            batch_embeddings = [item.embedding for item in response.data]
            all_embeddings.extend(batch_embeddings)

        return all_embeddings

    async def embed_single(self, text: str) -> List[float]:
        results = await self.embed_batch([text])
        return results[0]


# ──────────────────────────── 智谱 AI 实现 ────────────────────────────

class ZhipuEmbedding(EmbeddingProviderBase):
    """智谱 AI 嵌入模型"""

    def __init__(self) -> None:
        from zhipuai import ZhipuAI

        self._client = ZhipuAI(api_key=settings.ZHIPU_API_KEY)
        self._model = settings.EMBEDDING_MODEL or "embedding-3"

    @property
    def dimension(self) -> int:
        return settings.EMBEDDING_DIMENSIONS

    async def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """调用智谱 AI 接口批量生成嵌入"""
        import asyncio

        # 智谱 API 单次最多 16 条
        batch_size = 16
        all_embeddings: List[List[float]] = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            # 智谱 SDK 是同步的，用线程池包装
            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: self._client.embeddings.create(
                    input=batch,
                    model=self._model,
                ),
            )
            batch_embeddings = [item.embedding for item in response.data]
            all_embeddings.extend(batch_embeddings)

        return all_embeddings

    async def embed_single(self, text: str) -> List[float]:
        results = await self.embed_batch([text])
        return results[0]


# ──────────────────────────── 阿里云 DashScope 实现 ────────────────────────────

class DashscopeEmbedding(EmbeddingProviderBase):
    """阿里云 DashScope 嵌入模型"""

    def __init__(self) -> None:
        import dashscope
        dashscope.api_key = settings.DASHSCOPE_API_KEY
        self._model = settings.EMBEDDING_MODEL or "text-embedding-v3"

    @property
    def dimension(self) -> int:
        return settings.EMBEDDING_DIMENSIONS

    async def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """调用 DashScope 接口批量生成嵌入"""
        import asyncio
        from dashscope import TextEmbedding

        batch_size = 25  # DashScope 单次最多 25 条
        all_embeddings: List[List[float]] = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            loop = asyncio.get_event_loop()

            def _call() -> list:
                resp = TextEmbedding.call(
                    model=self._model,
                    input=batch,
                    dimension=self.dimension,
                )
                if resp.status_code != 200:
                    raise RuntimeError(f"DashScope 嵌入失败: {resp.code} - {resp.message}")
                return [item["embedding"] for item in resp.output["embeddings"]]

            batch_embeddings = await loop.run_in_executor(None, _call)
            all_embeddings.extend(batch_embeddings)

        return all_embeddings

    async def embed_single(self, text: str) -> List[float]:
        results = await self.embed_batch([text])
        return results[0]


# ──────────────────────────── 本地 BGE 实现 ────────────────────────────

class LocalBGEEmbedding(EmbeddingProviderBase):
    """
    本地 BGE 嵌入模型
    使用 sentence-transformers 加载模型，在 CPU/GPU 上本地推理
    """

    def __init__(self) -> None:
        self._model = None
        self._dim = settings.EMBEDDING_DIMENSIONS

    def _load_model(self):
        """延迟加载模型（首次调用时初始化）"""
        if self._model is not None:
            return

        logger.info("正在加载本地 BGE 嵌入模型: %s", settings.EMBEDDING_MODEL)
        from sentence_transformers import SentenceTransformer

        model_path = settings.LOCAL_BGE_MODEL_PATH or settings.EMBEDDING_MODEL
        self._model = SentenceTransformer(
            model_path,
            trust_remote_code=True,
        )
        # 获取实际维度
        test_embedding = self._model.encode(["测试"])
        self._dim = len(test_embedding[0])
        logger.info("BGE 模型加载完成，维度: %d", self._dim)

    @property
    def dimension(self) -> int:
        return self._dim

    async def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """使用本地模型批量生成嵌入"""
        import asyncio

        self._load_model()
        loop = asyncio.get_event_loop()

        def _encode() -> List[List[float]]:
            embeddings = self._model.encode(texts, normalize_embeddings=True)
            return embeddings.tolist()

        return await loop.run_in_executor(None, _encode)

    async def embed_single(self, text: str) -> List[float]:
        results = await self.embed_batch([text])
        return results[0]


# ──────────────────────────── MiniMax 实现 ────────────────────────────

class MiniMaxEmbedding(EmbeddingProviderBase):
    """
    MiniMax 嵌入模型
    通过 MiniMax 的 OpenAI 兼容接口调用嵌入 API
    """

    def __init__(self) -> None:
        from openai import AsyncOpenAI

        self._client = AsyncOpenAI(
            api_key=settings.MINIMAX_API_KEY or "EMPTY",
            base_url=settings.MINIMAX_BASE_URL,
        )
        self._model = settings.MINIMAX_EMBEDDING_MODEL

    @property
    def dimension(self) -> int:
        return settings.EMBEDDING_DIMENSIONS

    async def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """调用 MiniMax 兼容接口批量生成嵌入"""
        batch_size = 64
        all_embeddings: List[List[float]] = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            response = await self._client.embeddings.create(
                input=batch,
                model=self._model,
            )
            batch_embeddings = [item.embedding for item in response.data]
            all_embeddings.extend(batch_embeddings)

        return all_embeddings

    async def embed_single(self, text: str) -> List[float]:
        results = await self.embed_batch([text])
        return results[0]


# ──────────────────────────── 工厂类 ────────────────────────────

class EmbeddingService:
    """
    嵌入服务工厂类
    根据配置自动选择嵌入模型提供商，提供全局单例访问
    """

    _instance: Optional[EmbeddingProviderBase] = None

    @classmethod
    def get_instance(cls) -> EmbeddingProviderBase:
        """获取嵌入服务单例实例"""
        if cls._instance is not None:
            return cls._instance

        provider_map = {
            EmbeddingProvider.OPENAI: OpenAIEmbedding,
            EmbeddingProvider.ZHIPU: ZhipuEmbedding,
            EmbeddingProvider.DASHSCOPE: DashscopeEmbedding,
            EmbeddingProvider.LOCAL_BGE: LocalBGEEmbedding,
            EmbeddingProvider.MINIMAX: MiniMaxEmbedding,
        }

        provider_cls = provider_map.get(settings.EMBEDDING_PROVIDER)
        if provider_cls is None:
            raise ValueError(f"不支持的嵌入模型提供商: {settings.EMBEDDING_PROVIDER}")

        cls._instance = provider_cls()
        logger.info(
            "嵌入服务初始化完成: provider=%s, model=%s",
            settings.EMBEDDING_PROVIDER.value,
            settings.EMBEDDING_MODEL,
        )
        return cls._instance

    @classmethod
    async def embed_batch(cls, texts: List[str]) -> List[List[float]]:
        """便捷方法：批量生成嵌入"""
        instance = cls.get_instance()
        return await instance.embed_batch(texts)

    @classmethod
    async def embed_single(cls, text: str) -> List[float]:
        """便捷方法：生成单条嵌入"""
        instance = cls.get_instance()
        return await instance.embed_single(text)

    @classmethod
    def reset(cls) -> None:
        """重置单例（用于测试或切换配置）"""
        cls._instance = None