"""统一获取普通模型或绑定结构化输出的模型"""

from functools import cache
from typing import Literal, TypeVar, cast, overload

from langchain_core.language_models import LanguageModelInput
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.runnables import Runnable
from langchain_openai import ChatOpenAI
from pydantic import BaseModel

from app.config import settings

OutputModel = TypeVar("OutputModel", bound=BaseModel)
StructuredMethod = Literal["function_calling", "json_schema", "json_mode"]


@cache
def _get_base_llm() -> ChatOpenAI:
    """只缓存基础客户端，不因输出 schema 不同而重复创建连接"""
    if not settings.llm_api_key.strip() or not settings.llm_model.strip():
        raise ValueError("未配置 LLM_API_KEY 或 LLM_MODEL")
    return ChatOpenAI(
        model=settings.llm_model,
        api_key=settings.llm_api_key,
        base_url=settings.llm_base_url,
        timeout=settings.llm_timeout,
        max_retries=settings.llm_max_retries,
    )


@overload
def get_llm(
    schema: type[OutputModel],
    *,
    method: StructuredMethod = "function_calling",
    strict: bool | None = False,
    client: BaseChatModel | None = None,
) -> Runnable[LanguageModelInput, OutputModel]: ...


@overload
def get_llm(
    schema: None = None,
    *,
    method: StructuredMethod = "function_calling",
    strict: bool | None = False,
    client: None = None,
) -> ChatOpenAI: ...


@overload
def get_llm(
    schema: None = None,
    *,
    method: StructuredMethod = "function_calling",
    strict: bool | None = False,
    client: BaseChatModel,
) -> BaseChatModel: ...


def get_llm(
    schema: type[OutputModel] | None = None,
    *,
    method: StructuredMethod = "function_calling",
    strict: bool | None = False,
    client: BaseChatModel | None = None,
) -> BaseChatModel | Runnable[LanguageModelInput, OutputModel]:
    """不传 schema 返回普通模型，传入 Pydantic 类则返回结构化调用器

    client 用于注入模型，省略时使用缓存客户端
    strict 控制服务端 schema 约束，本地仍由 Pydantic 校验
    """
    if schema is not None and (
        not isinstance(schema, type) or not issubclass(schema, BaseModel)
    ):
        raise TypeError("schema 必须是 Pydantic BaseModel 类")
    if method not in ("function_calling", "json_schema", "json_mode"):
        raise ValueError(f"不支持的结构化输出方式: {method}")
    if schema is None and (method != "function_calling" or strict not in (False, None)):
        raise ValueError("设置结构化输出参数时必须提供 schema")
    if method == "json_mode" and strict is True:
        raise ValueError("json_mode 不支持 strict=True")

    llm = client if client is not None else _get_base_llm()
    if schema is None:
        return llm
    # json_mode 不接受 strict 参数，提示词还需要明确要求 JSON 及其结构
    structured = llm.with_structured_output(
        schema, method=method, strict=None if method == "json_mode" else strict
    )
    return cast(Runnable[LanguageModelInput, OutputModel], structured)