# app/module/shared/llm/factory.py
"""
LangChain 模型工厂：基于 OpenAI 兼容接口封装 Chat / Embedding 的创建、健康检查与同步调用。
- DashScope compatible‑mode：重写embedding逻辑，规避tiktoken token数组输入问题
- 其他OpenAI兼容服务商：完全使用原生 OpenAIEmbeddings / ChatOpenAI
"""
from typing import Literal, List, Optional

import httpx
from langchain_core.language_models import BaseChatModel
from langchain_core.embeddings import Embeddings
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from pydantic import SecretStr


class DashScopeOpenAIEmbeddings(OpenAIEmbeddings, Embeddings):
    """
    仅用于阿里云DashScope compatible‑mode
    1. 绕过原生tiktoken逻辑，直接传入原始字符串，不传入token数字数组
    2. 适配新版openai SDK返回Pydantic Embedding对象，使用属性 .embedding
    """

    def _get_len_safe_embeddings(self, texts: List[str], *, engine: str, **kwargs) -> List[List[float]]:
        batch_size = self.chunk_size
        embeddings: List[List[float]] = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            resp = self.client.create(
                input=batch,
                model=self.model,
                encoding_format="float",
            )
            embeddings.extend([data.embedding for data in resp.data])
        return embeddings


class LLMFactory:
    @staticmethod
    def _is_dashscope_compatible(base_url: Optional[str]) -> bool:
        if not base_url:
            return False
        return "dashscope.aliyuncs.com/compatible-mode" in base_url

    @staticmethod
    def create_chat(
        model_name: str,
        base_url: str,
        api_key: str | None = None,
    ) -> BaseChatModel:
        """创建对话模型，兼容 OpenAI 及任意 base_url 的 OpenAI 兼容服务。"""
        # DashScope容器环境关闭SSL校验；其余服务商使用默认httpx客户端（开启证书校验）
        if LLMFactory._is_dashscope_compatible(base_url):
            http_client = httpx.Client(verify=False)
        else:
            http_client = None

        return ChatOpenAI(
            model=model_name,
            base_url=base_url or None,
            api_key=SecretStr(api_key or ""),
            http_client=http_client,
        )

    @staticmethod
    def create_embedding(
        model_name: str,
        base_url: str,
        api_key: str | None = None,
    ) -> Embeddings:
        is_dashscope = LLMFactory._is_dashscope_compatible(base_url)

        # DashScope才关闭SSL校验，其他服务商走默认SSL校验
        if is_dashscope:
            http_client = httpx.Client(verify=False)
        else:
            http_client = None

        kwargs = {
            "model": model_name,
            "base_url": base_url or None,
            "api_key": SecretStr(api_key or ""),
            "http_client": http_client,
        }

        # 仅DashScope注入encoding_format="float"，其他服务商不强制注入，避免参数不兼容
        if is_dashscope:
            import inspect
            sig = inspect.signature(OpenAIEmbeddings.__init__)
            if "encoding_format" in sig.parameters:
                kwargs["encoding_format"] = "float"
            return DashScopeOpenAIEmbeddings(**kwargs)

        # 其余所有OpenAI兼容模型，走原生实现
        return OpenAIEmbeddings(**kwargs)

    @staticmethod
    def check_health(
        model_name: str,
        base_url: str,
        api_key: str | None,
        model_type: Literal["CHAT", "EMBEDDING"] | str,
    ) -> None:
        """对配置做一次最小化调用进行健康检查，失败则抛出。"""
        if model_type == "CHAT":
            model = LLMFactory.create_chat(model_name, base_url, api_key)
            model.invoke("hello")
        else:
            model = LLMFactory.create_embedding(model_name, base_url, api_key)
            model.embed_query("text")

    @staticmethod
    def get_embedding_dimension(
        model_name: str,
        base_url: str,
        api_key: str | None = None,
    ) -> int:
        """创建 Embedding 模型并返回向量维度。"""
        emb = LLMFactory.create_embedding(model_name, base_url, api_key)
        return len(emb.embed_query("text"))

    @staticmethod
    def invoke_sync(chat_model: BaseChatModel, prompt: str) -> str:
        """同步调用对话模型并返回 content，供 run_in_executor 等场景使用。"""
        return chat_model.invoke(prompt).content