from __future__ import annotations

from typing import TYPE_CHECKING, Iterator

from niannian.llm.prompts import load_prompt, render_prompt
from niannian.llm.types import LLMUnavailable, Message
from niannian.personas import RIGHT
from niannian.schemas.models import Atom, WikiHit

if TYPE_CHECKING:
    from niannian.llm.router import LLMRouter


class RightReason:
    """小思 — 逻辑整合；不持久存储，只组装当次输出。"""

    def answer(
        self,
        question: str,
        *,
        wiki_hits: list[WikiHit],
        summaries: list[str],
        atoms: list[Atom],
        llm: LLMRouter | None = None,
    ) -> str:
        if llm is not None and llm.is_module_configured("right"):
            try:
                return self.answer_llm(
                    question,
                    wiki_hits=wiki_hits,
                    summaries=summaries,
                    atoms=atoms,
                    llm=llm,
                )
            except LLMUnavailable:
                pass
        return self._rule_answer(question, wiki_hits=wiki_hits, summaries=summaries, atoms=atoms)

    def answer_llm(
        self,
        question: str,
        *,
        wiki_hits: list[WikiHit],
        summaries: list[str],
        atoms: list[Atom],
        llm: LLMRouter,
    ) -> str:
        memory_block = self.build_context_block(
            question,
            wiki_hits=wiki_hits,
            summaries=summaries,
            atoms=atoms,
        )
        system = load_prompt("right_query").get("system", "")
        prompt = render_prompt(
            "right_reason",
            question=question,
            memory_block=memory_block,
        )
        result = llm.complete(
            "right",
            [
                Message(role="system", content=system or prompt),
                Message(role="user", content=question if system else prompt),
            ],
        )
        return result.text

    def _rule_answer(
        self,
        question: str,
        *,
        wiki_hits: list[WikiHit],
        summaries: list[str],
        atoms: list[Atom],
    ) -> str:
        parts: list[str] = []
        if wiki_hits:
            parts.append("关联概念：" + "、".join(h.concept for h in wiki_hits[:3]))
        if summaries:
            parts.append("提炼摘要：" + "；".join(summaries[:2]))
        if atoms:
            parts.append("近期碎片：" + atoms[-1].raw_text[:160])
        if not parts:
            return f"（{RIGHT}）暂未找到与「{question}」相关的记忆，可先让小忆记下碎片。"
        return f"（{RIGHT}）基于记忆联想：{' '.join(parts)}"

    def build_context_block(
        self,
        task: str,
        *,
        wiki_hits: list[WikiHit],
        summaries: list[str],
        atoms: list[Atom],
    ) -> str:
        lines = [f"## 念念辅助上下文（任务：{task}）", ""]
        if wiki_hits:
            lines.append("### 联想（小织）")
            for h in wiki_hits[:5]:
                links = " → ".join(h.links[:4]) if h.links else ""
                lines.append(f"- [[{h.concept}]] {h.summary[:100]} {links}")
            lines.append("")
        if summaries:
            lines.append("### 提炼（小海）")
            for s in summaries[:3]:
                lines.append(f"- {s}")
            lines.append("")
        if atoms:
            lines.append("### 碎片（小忆）")
            for a in atoms[-3:]:
                tag_str = ",".join(a.tags) if a.tags else "无标签"
                lines.append(f"- [{tag_str}] {a.raw_text[:120]}")
            lines.append("")
        return "\n".join(lines).strip()

    def build_chat_messages(
        self,
        user_text: str,
        *,
        memory_block: str,
        history: list[Message],
        skill_block: str = "",
        task_block: str = "",
        subject_block: str = "",
        object_block: str = "",
        persona_block: str = "",
        character_block: str = "",
    ) -> list[Message]:
        """Prompt 四层栈（chat-llm-design §5）+ v0.5 主客体双块 + v0.6 persona overlay + v0.12 VP。"""
        system_persona = load_prompt("right_chat").get("system", "").strip()
        memory_section = memory_block.strip() or "（本轮无相关记忆）"
        messages: list[Message] = [
            Message(role="system", content=system_persona),
        ]
        if subject_block.strip():
            messages.append(Message(role="system", content=subject_block.strip()))
        if character_block.strip():
            messages.append(Message(role="system", content=character_block.strip()))
        if persona_block.strip():
            messages.append(Message(role="system", content=persona_block.strip()))
        if object_block.strip():
            messages.append(Message(role="system", content=object_block.strip()))
        messages.append(
            Message(
                role="system",
                content=f"## 念念记忆包（本轮召回）\n{memory_section}",
            ),
        )
        if skill_block.strip():
            messages.append(
                Message(role="system", content=f"## 念念 Skill 注入\n{skill_block.strip()}"),
            )
        if task_block.strip():
            messages.append(
                Message(role="system", content=f"## 当前任务\n{task_block.strip()}"),
            )
        messages.extend(history)
        messages.append(Message(role="user", content=user_text))
        return messages

    def complete_chat(
        self,
        messages: list[Message],
        llm: LLMRouter,
        *,
        stream: bool = False,
    ) -> str | Iterator[str]:
        if stream:
            return self._stream_chat(messages, llm)
        return llm.complete("right", messages).text

    def _stream_chat(self, messages: list[Message], llm: LLMRouter) -> Iterator[str]:
        for chunk in llm.stream_complete("right", messages):
            if chunk.text:
                yield chunk.text

    @staticmethod
    def fallback_chat_reply(messages: list[Message], *, memory_block: str = "") -> str:
        user = next((m.content for m in reversed(messages) if m.role == "user"), "")
        if memory_block:
            return (
                f"（{RIGHT}）云端 LLM 暂不可用。以下是召回的记忆摘要：\n\n"
                f"{memory_block[:600]}\n\n"
                f"收到你的消息：「{user[:120]}」"
            )
        return f"（{RIGHT}）你好！已收到「{user[:120]}」。当前未配置云端 LLM，可使用 query/assist 接口。"

    @staticmethod
    def trace_line(msg: str) -> str:
        return f"[{RIGHT}] {msg}"

