from __future__ import annotations

from typing import Any, TypedDict

from langgraph.graph import END, START, StateGraph

from app.schemas.steward import (
    StewardRuntimeDecisionRequest,
    StewardRuntimeDecisionResponse,
    StewardSlotDecisionRequest,
    StewardSlotDecisionResponse,
)
from app.services.runtime_chat import RuntimeChatService
from app.services.steward_runtime_decision_agent import StewardRuntimeDecisionAgent
from app.services.steward_slot_decision_agent import StewardSlotDecisionAgent


class StewardSlotGraphState(TypedDict, total=False):
    request: StewardSlotDecisionRequest
    normalized_request: StewardSlotDecisionRequest
    decision: StewardSlotDecisionResponse
    model_call_traces: list[dict[str, Any]]
    fallback_reason: str


class StewardRuntimeGraphState(TypedDict, total=False):
    request: StewardRuntimeDecisionRequest
    normalized_request: StewardRuntimeDecisionRequest
    action_decision: StewardRuntimeDecisionResponse
    decision: StewardRuntimeDecisionResponse
    model_call_traces: list[dict[str, Any]]
    fallback_reason: str


class StewardGraphRuntime:
    """用 LangGraph 编排会话内槽位、记忆和行动决策。"""

    def __init__(self, runtime_chat_service: RuntimeChatService) -> None:
        self.slot_agent = StewardSlotDecisionAgent(runtime_chat_service)
        self.runtime_agent = StewardRuntimeDecisionAgent(runtime_chat_service)
        self._slot_graph = self._build_slot_graph()
        self._runtime_graph = self._build_runtime_graph()

    def decide_slot(self, request: StewardSlotDecisionRequest) -> StewardSlotDecisionResponse:
        final_state = self._slot_graph.invoke(
            {
                "request": request,
                "model_call_traces": [],
                "fallback_reason": "",
            }
        )
        decision = final_state.get("decision")
        if not isinstance(decision, StewardSlotDecisionResponse):
            raise RuntimeError("LangGraph 槽位决策未生成有效结果。")
        return decision

    def decide_runtime(self, request: StewardRuntimeDecisionRequest) -> StewardRuntimeDecisionResponse:
        final_state = self._runtime_graph.invoke(
            {
                "request": request,
                "model_call_traces": [],
                "fallback_reason": "",
            }
        )
        decision = final_state.get("decision")
        if not isinstance(decision, StewardRuntimeDecisionResponse):
            raise RuntimeError("LangGraph 运行时决策未生成有效结果。")
        return decision

    def _build_slot_graph(self):
        graph = StateGraph(StewardSlotGraphState)
        graph.add_node("slot_prepare_context", self._slot_prepare_context)
        graph.add_node("slot_tool_decision", self._slot_tool_decision)
        graph.add_node("slot_rule_fallback", self._slot_rule_fallback)

        graph.add_edge(START, "slot_prepare_context")
        graph.add_edge("slot_prepare_context", "slot_tool_decision")
        graph.add_conditional_edges(
            "slot_tool_decision",
            self._route_after_slot_tool_decision,
            {
                "done": END,
                "fallback": "slot_rule_fallback",
            },
        )
        graph.add_edge("slot_rule_fallback", END)
        return graph.compile()

    def _build_runtime_graph(self):
        graph = StateGraph(StewardRuntimeGraphState)
        graph.add_node("runtime_memory_context", self._runtime_memory_context)
        graph.add_node("runtime_action_decision", self._runtime_action_decision)
        graph.add_node("runtime_tool_decision", self._runtime_tool_decision)
        graph.add_node("runtime_rule_fallback", self._runtime_rule_fallback)

        graph.add_edge(START, "runtime_memory_context")
        graph.add_conditional_edges(
            "runtime_memory_context",
            self._route_after_runtime_memory_context,
            {
                "action": "runtime_action_decision",
                "tool": "runtime_tool_decision",
            },
        )
        graph.add_edge("runtime_action_decision", END)
        graph.add_conditional_edges(
            "runtime_tool_decision",
            self._route_after_runtime_tool_decision,
            {
                "done": END,
                "fallback": "runtime_rule_fallback",
            },
        )
        graph.add_edge("runtime_rule_fallback", END)
        return graph.compile()

    def _slot_prepare_context(self, state: StewardSlotGraphState) -> dict[str, Any]:
        return {
            "normalized_request": self.slot_agent._normalize_request(state["request"]),
        }

    def _slot_tool_decision(self, state: StewardSlotGraphState) -> dict[str, Any]:
        request = state.get("normalized_request") or state["request"]
        try:
            return {"decision": self.slot_agent.decide(request)}
        except Exception as exc:
            return {
                "model_call_traces": [
                    *state.get("model_call_traces", []),
                    self._build_failure_trace("langgraph_slot_decision", exc),
                ],
                "fallback_reason": f"LangGraph 槽位工具节点失败，已切换规则兜底：{exc}",
            }

    @staticmethod
    def _route_after_slot_tool_decision(state: StewardSlotGraphState) -> str:
        if isinstance(state.get("decision"), StewardSlotDecisionResponse):
            return "done"
        return "fallback"

    def _slot_rule_fallback(self, state: StewardSlotGraphState) -> dict[str, StewardSlotDecisionResponse]:
        request = state.get("normalized_request") or self.slot_agent._normalize_request(state["request"])
        return {
            "decision": self.slot_agent._build_rule_fallback(
                request,
                state.get("model_call_traces", []),
            )
        }

    def _runtime_memory_context(self, state: StewardRuntimeGraphState) -> dict[str, Any]:
        request = self.runtime_agent._normalize_request(state["request"])
        action_decision = self.runtime_agent._build_selected_flow_decision(request, [])
        update: dict[str, Any] = {"normalized_request": request}
        if action_decision is not None:
            update["action_decision"] = action_decision
        return update

    @staticmethod
    def _route_after_runtime_memory_context(state: StewardRuntimeGraphState) -> str:
        if isinstance(state.get("action_decision"), StewardRuntimeDecisionResponse):
            return "action"
        return "tool"

    def _runtime_action_decision(
        self,
        state: StewardRuntimeGraphState,
    ) -> dict[str, StewardRuntimeDecisionResponse]:
        return {"decision": state["action_decision"]}

    def _runtime_tool_decision(self, state: StewardRuntimeGraphState) -> dict[str, Any]:
        request = state.get("normalized_request") or state["request"]
        try:
            return {"decision": self.runtime_agent.decide(request)}
        except Exception as exc:
            return {
                "model_call_traces": [
                    *state.get("model_call_traces", []),
                    self._build_failure_trace("langgraph_runtime_decision", exc),
                ],
                "fallback_reason": f"LangGraph 运行时工具节点失败，已切换规则兜底：{exc}",
            }

    @staticmethod
    def _route_after_runtime_tool_decision(state: StewardRuntimeGraphState) -> str:
        if isinstance(state.get("decision"), StewardRuntimeDecisionResponse):
            return "done"
        return "fallback"

    def _runtime_rule_fallback(
        self,
        state: StewardRuntimeGraphState,
    ) -> dict[str, StewardRuntimeDecisionResponse]:
        request = state.get("normalized_request") or self.runtime_agent._normalize_request(state["request"])
        decision = self.runtime_agent._build_rule_fallback(
            request,
            state.get("model_call_traces", []),
        )
        return {
            "decision": self.runtime_agent._attach_updated_steward_state(
                decision,
                request,
            )
        }

    @staticmethod
    def _build_failure_trace(slot: str, exc: Exception) -> dict[str, Any]:
        return {
            "slot": slot,
            "provider": "langgraph",
            "model": "",
            "attempt": 1,
            "status": "failed",
            "error": str(exc),
        }