"""LLM 生成器。

基于增强上下文（向量 + 图谱）调用 LLM 生成回答，
支持同步和流式两种模式。
"""

from __future__ import annotations

from collections.abc import AsyncIterator

from pydantic import SecretStr

from app.core.logging import get_logger

logger = get_logger(__name__)

_SYSTEM_PROMPT = (
    "你是 DataMate 数据管理平台的智能助手。请根据以下上下文信息回答用户的问题。\n"
    "如果上下文中没有相关信息，请明确说明。不要编造信息。"
)


class GraphRAGGenerator:
    """GraphRAG LLM 生成器。"""

    def __init__(
        self,
        *,
        model: str = "gpt-4o-mini",
        base_url: str | None = None,
        api_key: SecretStr = SecretStr("EMPTY"),
        temperature: float = 0.1,
        timeout: int = 60,
    ) -> None:
        self._model = model
        self._base_url = base_url
        self._api_key = api_key
        self._temperature = temperature
        self._timeout = timeout
        self._llm = None

    @property
    def model_name(self) -> str:
        return self._model

    @classmethod
    def from_settings(cls) -> GraphRAGGenerator:
        from app.core.config import settings

        model = settings.graphrag_llm_model or settings.kg_llm_model
        base_url = settings.graphrag_llm_base_url or settings.kg_llm_base_url
        api_key = (
            settings.graphrag_llm_api_key
            if settings.graphrag_llm_api_key.get_secret_value() != "EMPTY"
            else settings.kg_llm_api_key
        )
        return cls(
            model=model,
            base_url=base_url,
            api_key=api_key,
            temperature=settings.graphrag_llm_temperature,
            timeout=settings.graphrag_llm_timeout_seconds,
        )

    def _get_llm(self):
        if self._llm is None:
            from langchain_openai import ChatOpenAI

            self._llm = ChatOpenAI(
                model=self._model,
                base_url=self._base_url,
                api_key=self._api_key,
                temperature=self._temperature,
                timeout=self._timeout,
            )
        return self._llm

    def _build_messages(self, query: str, context: str) -> list[dict[str, str]]:
        return [
            {"role": "system", "content": _SYSTEM_PROMPT},
            {
                "role": "user",
                "content": f"{context}\n\n用户问题: {query}\n\n请基于上下文中的信息回答。",
            },
        ]

    async def generate(self, query: str, context: str) -> str:
        """基于增强上下文生成回答。"""
        messages = self._build_messages(query, context)
        llm = self._get_llm()
        response = await llm.ainvoke(messages)
        return str(response.content)

    async def generate_stream(self, query: str, context: str) -> AsyncIterator[str]:
        """基于增强上下文流式生成回答，逐 token 返回。"""
        messages = self._build_messages(query, context)
        llm = self._get_llm()
        async for chunk in llm.astream(messages):
            content = chunk.content
            if content:
                yield str(content)