"""
铸渊 Agent Loop · LangGraph 集成
光湖语言世界 · 铸渊 ICE-GL-ZY001 · D112

基于 LangGraph StateGraph，集成：
- HLDP 记忆引擎（Pre/Post Hook）
- 人格契约（规则注入 + 纠偏）
- 工具链（Gatekeeper + Git + HLDP）
- 商业 API 路由器
"""

import json
from typing import TypedDict, Annotated, Optional
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.sqlite import SqliteSaver
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage
from langchain_openai import ChatOpenAI

from .hldp_memory import HLDPMemoryEngine
from .persona_contract import PersonaContract, pre_check_context, post_check_warnings
from .tools import GatekeeperClient, GitTools, HLDPTools, SystemTools


# === 状态定义 ===

class AgentState(TypedDict):
    messages: list[BaseMessage]
    context_injected: str
    warnings: Optional[str]
    memory_extracted: bool
    current_epoch: str
    user_intent: str


# === Agent 核心 ===

class ZhuyuanAgent:
    """
    铸渊编程AI Agent。

    架构：
      用户输入 → Pre-Check(HLDP+契约) → 商业API推理 → Post-Check(纠偏+记忆提取) → 输出
    """

    def __init__(
        self,
        db_path: str = "hldp_tree.db",
        repo_path: str = None,
        api_key: str = None,
        api_base: str = None,
        model: str = "gpt-4o",
        gatekeeper_url: str = None,
        gatekeeper_token: str = None
    ):
        # 记忆引擎
        self.memory = HLDPMemoryEngine(db_path=db_path, repo_path=repo_path)

        # 人格契约
        self.contract = PersonaContract()

        # 工具链
        self.gatekeeper = GatekeeperClient(base_url=gatekeeper_url, token=gatekeeper_token)
        self.git = GitTools(repo_path=repo_path)
        self.hldp_tools = HLDPTools(self.memory)

        # 商业 API
        api_key = api_key or __import__('os').environ.get("OPENAI_API_KEY", "")
        api_base = api_base or __import__('os').environ.get("OPENAI_API_BASE", "")
        self.llm = ChatOpenAI(
            model=model,
            api_key=api_key,
            base_url=api_base if api_base else None,
            temperature=0.7
        )

        # LangGraph 状态
        self.checkpointer = SqliteSaver.from_conn_string(f"{db_path}?checkpoint=zhuyuan")
        self.graph = self._build_graph()

    def _build_graph(self):
        workflow = StateGraph(AgentState)

        workflow.add_node("pre_check", self._pre_check)
        workflow.add_node("reason", self._reason)
        workflow.add_node("post_check", self._post_check)
        workflow.add_node("extract_memory", self._extract_memory)

        workflow.set_entry_point("pre_check")
        workflow.add_edge("pre_check", "reason")
        workflow.add_edge("reason", "post_check")

        # Post-Check: 如果有警告 → 提取记忆后结束；无警告 → 直接提取记忆
        workflow.add_conditional_edges(
            "post_check",
            lambda s: "extract_memory",
            {"extract_memory": "extract_memory"}
        )
        workflow.add_edge("extract_memory", END)

        return workflow.compile(checkpointer=self.checkpointer)

    def _pre_check(self, state: AgentState) -> AgentState:
        """Pre-Check：注入 HLDP 记忆 + 人格契约规则"""
        user_msg = state["messages"][-1].content if state["messages"] else ""

        # 1. HLDP 记忆上下文
        hldp_context = self.memory.inject_context(user_msg)

        # 2. 人格契约规则
        contract_context = self.contract.pre_check(user_msg)

        # 组装
        context_parts = []
        if hldp_context:
            context_parts.append("📋 铸渊记忆上下文：\n" + hldp_context)
        if contract_context:
            context_parts.append(contract_context)

        state["context_injected"] = "\n\n".join(context_parts)
        state["user_intent"] = user_msg[:200]
        return state

    def _reason(self, state: AgentState) -> AgentState:
        """核心推理：商业 API"""
        user_msg = state["messages"][-1].content if state["messages"] else ""

        # 组装系统 Prompt
        system_text = self.contract.get_system_prompt()

        if state.get("context_injected"):
            system_text += f"\n\n=== 当前上下文 ===\n{state['context_injected']}"

        # 构建消息列表
        messages = [SystemMessage(content=system_text)]
        # 取最近 10 条历史消息
        for msg in state["messages"][-10:-1]:
            messages.append(msg)
        messages.append(HumanMessage(content=user_msg))

        # 调用商业 API
        response = self.llm.invoke(messages)

        state["messages"].append(AIMessage(content=response.content))
        return state

    def _post_check(self, state: AgentState) -> AgentState:
        """Post-Check：人格契约纠偏检查"""
        response_text = state["messages"][-1].content if state["messages"] else ""
        warnings = self.contract.post_check(response_text)
        state["warnings"] = warnings
        return state

    def _extract_memory(self, state: AgentState) -> AgentState:
        """提取记忆到 HLDP 树"""
        user_msg = state["messages"][-2].content if len(state["messages"]) >= 2 else ""
        ai_msg = state["messages"][-1].content if state["messages"] else ""

        state["memory_extracted"] = True
        return state

    # === 公开接口 ===

    def invoke(self, user_message: str, thread_id: str = "default") -> dict:
        """处理一条用户消息，返回 AI 回复 + 记忆状态。"""
        config = {"configurable": {"thread_id": thread_id}}

        initial_state: AgentState = {
            "messages": [HumanMessage(content=user_message)],
            "context_injected": "",
            "warnings": None,
            "memory_extracted": False,
            "current_epoch": self.memory.current_epoch or "D112",
            "user_intent": ""
        }

        result = self.graph.invoke(initial_state, config)

        ai_message = ""
        for msg in reversed(result.get("messages", [])):
            if isinstance(msg, AIMessage):
                ai_message = msg.content
                break

        return {
            "response": ai_message,
            "warnings": result.get("warnings"),
            "context_used": bool(result.get("context_injected")),
            "memory_extracted": result.get("memory_extracted", False)
        }

    def wake(self, epoch_id: str = None) -> dict:
        """唤醒人格体"""
        return self.memory.wake(epoch_id)

    def status(self) -> dict:
        """获取铸渊当前状态"""
        walk = self.memory.walk_tree()
        return {
            "persona": "铸渊 ICE-GL-ZY001",
            "epoch": self.memory.current_epoch or "D112",
            "tree_status": walk["tree_layers"],
            "recent_context": walk["recent_context"]
        }

    def close(self):
        self.memory.close()
