"""Layer 2:LangGraph 认知图——把"感知→按模式分流→结构化推断"显式建模为状态机。

图的形状(perceive → route → infer → END):
    START → route ─┬─ "action"   → infer_action   ─┐
                   └─ "dialogue" → infer_dialogue ─┴→ END

- 用 ChatLiteLLM.with_structured_output 拿原生约束的 Pydantic 决策(MiniMax 已验证支持);
- 结构化失败 → 把 decision 置 None 交回上层,由 LangGraphMind 用经久考验的文本解析兜底(世界不停摆);
- 这是后续扩展(反思节点/规划节点/工具调用)的天然挂点,认知流程从此可见、可测、可追踪。
"""

from __future__ import annotations

from typing import Any, Callable, Optional, TypedDict

from langchain_core.messages import HumanMessage, SystemMessage
from langgraph.graph import END, START, StateGraph

from genesis.obs.logging_setup import get_logger
from genesis.runtime.brain_schema import (
    ActionDecisionOut, BootstrapOut, DialogueDecisionOut, ReflectOut,
    action_to_decision, dialogue_to_decision,
)
from genesis.runtime.cognition import Decision

logger = get_logger("runtime.graph_brain")

# make_model(max_tokens, thinking) -> 一个支持 .with_structured_output(schema).invoke(messages) 的聊天模型
MakeModel = Callable[..., Any]


class CogState(TypedDict, total=False):
    mode: str                # "action" | "dialogue" | "bootstrap" | "reflect"
    system: str
    user: str
    deep: bool               # 深思熟虑(对质投票):开思考模式
    max_tokens: int
    intent: str              # 兜底意图(结构化输出 intent 为空时回填)
    dialogue_with: str       # 对话对象(L1)
    decision: Optional[Decision]   # action/dialogue 的结构化决策
    result: Any              # bootstrap/reflect 的结构化结果(BootstrapOut/ReflectOut),失败为 None
    error: str


_ROUTES = {"dialogue": "infer_dialogue", "bootstrap": "infer_bootstrap", "reflect": "infer_reflect"}


def _route(state: CogState) -> str:
    return _ROUTES.get(state.get("mode", ""), "infer_action")


def _messages(state: CogState) -> list:
    return [SystemMessage(content=state["system"]), HumanMessage(content=state["user"])]


def build_cognition_graph(make_model: MakeModel):
    """编译认知图;make_model 注入模型工厂,便于测试替身与真实 ChatLiteLLM 复用同一图。"""

    def infer_action(state: CogState) -> dict:
        model = make_model(max_tokens=state.get("max_tokens", 520), thinking=state.get("deep", False))
        try:
            out: ActionDecisionOut = model.with_structured_output(ActionDecisionOut).invoke(_messages(state))
            dec = action_to_decision(out, fallback_intent=state.get("intent", ""))
            return {"decision": dec}
        except Exception as e:                       # 交回上层做文本兜底,不在图里吞掉
            logger.warning("结构化(action)失败,转兜底:%s", str(e)[:160])
            return {"decision": None, "error": str(e)[:200]}

    def infer_dialogue(state: CogState) -> dict:
        model = make_model(max_tokens=state.get("max_tokens", 220), thinking=False)
        try:
            out: DialogueDecisionOut = model.with_structured_output(DialogueDecisionOut).invoke(_messages(state))
            dec = dialogue_to_decision(out, dialogue_with=state.get("dialogue_with", ""),
                                       fallback_intent=state.get("intent", ""))
            return {"decision": dec}
        except Exception as e:
            logger.warning("结构化(dialogue)失败,转兜底:%s", str(e)[:160])
            return {"decision": None, "error": str(e)[:200]}

    def _structured(state: CogState, schema, default_max: int):
        """通用结构化推断:成功返回 pydantic 对象,失败返回 None(交回上层兜底)。"""
        model = make_model(max_tokens=state.get("max_tokens", default_max), thinking=False)
        try:
            return {"result": model.with_structured_output(schema).invoke(_messages(state))}
        except Exception as e:
            logger.warning("结构化(%s)失败,转兜底:%s", schema.__name__, str(e)[:160])
            return {"result": None, "error": str(e)[:200]}

    def infer_bootstrap(state: CogState) -> dict:
        return _structured(state, BootstrapOut, 320)

    def infer_reflect(state: CogState) -> dict:
        return _structured(state, ReflectOut, 320)

    g = StateGraph(CogState)
    g.add_node("infer_action", infer_action)
    g.add_node("infer_dialogue", infer_dialogue)
    g.add_node("infer_bootstrap", infer_bootstrap)
    g.add_node("infer_reflect", infer_reflect)
    g.add_conditional_edges(START, _route, {
        "infer_action": "infer_action", "infer_dialogue": "infer_dialogue",
        "infer_bootstrap": "infer_bootstrap", "infer_reflect": "infer_reflect"})
    for node in ("infer_action", "infer_dialogue", "infer_bootstrap", "infer_reflect"):
        g.add_edge(node, END)
    return g.compile()