from collections.abc import AsyncIterator

from app.agents.llm_adapter import LLMResponse, LLMStreamChunk, OpenAICompatibleLLMClient
from app.core.config import settings
from app.schemas.learning_assistant import LearningAssistantSource


class LearningAssistantAgent:
    """AI学习助手 Agent：根据 RAG 来源生成带循证出处的医学学习回答。"""

    def __init__(self, llm_client: OpenAICompatibleLLMClient | None = None) -> None:
        self.llm_client = llm_client or OpenAICompatibleLLMClient()

    async def answer(self, question: str, sources: list[LearningAssistantSource]) -> LLMResponse:
        """非流式回答：把问题和检索来源拼接后调用快速模型生成标准回答。"""
        return await self.llm_client.chat(
            self._messages(question, sources),
            model=settings.llm_fast_model,
            thinking_enabled=settings.llm_fast_thinking_enabled,
            max_tokens=1200,
        )

    async def stream_answer(self, question: str, sources: list[LearningAssistantSource]) -> AsyncIterator[LLMStreamChunk]:
        """流式回答：输出 AI 学习助手增量文本，前端可直接渲染。"""
        async for chunk in self.llm_client.stream_chat(
            self._messages(question, sources),
            model=settings.llm_fast_model,
            thinking_enabled=settings.llm_fast_thinking_enabled,
            max_tokens=1200,
        ):
            yield chunk

    def _messages(self, question: str, sources: list[LearningAssistantSource]) -> list[dict]:
        """提示词拼接：命中知识库时必须引用来源，未命中时必须声明未找到参考。"""
        if sources:
            context = "\n\n".join(
                (
                    f"[来源{index}] 文档：{source.document_title or source.file_name}；"
                    f"页码：{source.page_start}-{source.page_end}；chunk_uid：{source.chunk_uid}\n"
                    f"{source.quote}"
                )
                for index, source in enumerate(sources, start=1)
            )
            system = (
                "你是医学学习助手，只用于医学教育学习，不替代临床诊疗。"
                "请优先依据给定知识库片段回答，回答要清晰、准确、分点。"
                "每个关键结论后标注对应来源编号，例如【来源1】。"
                "不得编造不存在的PDF、页码或指南来源。"
            )
            user = f"用户问题：{question}\n\n可用知识库片段：\n{context}\n\n请给出带来源的学习回答。"
        else:
            system = (
                "你是医学学习助手，只用于医学教育学习，不替代临床诊疗。"
                "当前没有检索到机构知识库参考，回答开头必须写：未检索到本机构知识库参考，以下为大模型通用学习回答。"
                "不得伪造PDF来源、页码或指南名称。"
            )
            user = f"用户问题：{question}\n\n请给出通用学习回答，并提醒用户以课程教材和临床规范为准。"
        return [{"role": "system", "content": system}, {"role": "user", "content": user}]