from langchain.agents import create_agent
from langchain.agents.middleware import (
    before_model, after_model
)
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_core.tools import tool

# ---- 定义工具 ----
@tool
def get_news(query: str) -> str:
    """查询新闻"""
    return f"关于 {query} 的最新新闻：今日AI领域有重大突破..."

# ---- 钩子函数1: 伪造 AIMessage，跳过模型调用直接到工具 ----
@before_model(can_jump_to="tools")
def fake_ai_message(state, runtime):
    """检测到 direct_to_tools 字段时，伪造 AIMessage 跳过模型调用"""
    messages = state["messages"]
    last_msg = messages[-1]
    if hasattr(last_msg, "content") and "direct_to_tools" in last_msg.content:
        # 伪造一条 AIMessage，包含工具调用
        return {
            "messages": [AIMessage(
                content="",
                tool_calls=[{
                    "name": "get_news",
                    "args": {"query": "AI新闻"},
                    "id": "call_001"
                }]
            )]
        }
    return None

# ---- 钩子函数2: 模型调用后不满意，强制重新调用 ----
@after_model(can_jump_to="model")
def retry_model(state, runtime):
    """检测到 retry_model 字段且未重试过时，追加 SystemMessage 重新调用"""
    messages = state["messages"]
    # 检查用户原始输入是否包含 retry_model
    user_msgs = [m for m in messages if isinstance(m, HumanMessage)]
    if not user_msgs:
        return None
    last_user = user_msgs[-1]
    if "retry_model" not in last_user.content:
        return None
    # 检查是否已经重试过（避免死循环）
    ai_msgs = [m for m in messages if isinstance(m, AIMessage)]
    if ai_msgs and "二次回答" in ai_msgs[-1].content:
        return None  # 已经重试过，不再重复
    # 追加 SystemMessage 要求重试
    return {
        "messages": [SystemMessage(
            content="你必须以'二次回答'开头，并且只用一句话回答。"
        )]
    }

# ---- 钩子函数3: 溢出检测，直接结束 ----
@before_model(can_jump_to="end")
def check_overflow(state, runtime):
    """检测到 overflow 字段时，提前终止 Agent"""
    messages = state["messages"]
    last_msg = messages[-1]
    if hasattr(last_msg, "content") and "overflow" in last_msg.content:
        return {
            "messages": [AIMessage(content="上下文窗口溢出，已终止执行。")]
        }
    return None

# ---- 创建 Agent ----
agent = create_agent(
    model=model,
    tools=[get_news],
    middleware=[
        fake_ai_message,    # before_model: 跳 tools
        retry_model,        # after_model: 跳 model
        check_overflow      # before_model: 跳 end
    ]
)

# ---- 测试用例 ----
def test(query: str):
    print(f"\n{'='*50}")
    print(f"输入: {query}")
    response = agent.invoke({"messages": [HumanMessage(content=query)]})
    for msg in response["messages"]:
        msg.pretty_print()

# Case 1: 跳过模型，直接到工具
test("请帮我查询今日天气 direct_to_tools")

# Case 2: 模型调用后重试
test("请随便介绍一下LangChain retry_model")

# Case 3: 溢出直接结束
test("你好 overflow")

# Case 4: 正常执行
test("今天有什么新闻？")

from langchain.agents.middleware import AgentMiddleware, hook_configure

class MyMiddleware(AgentMiddleware):

    @hook_configure(can_jump_to="tools")
    @hook_configure(can_jump_to="end")
    def before_model(self, state, runtime):
        """合并两个 before_model 跳转逻辑"""
        messages = state["messages"]
        last_msg = messages[-1]
        content = last_msg.content if hasattr(last_msg, "content") else ""

        # 分支1: 跳转到工具
        if "direct_to_tools" in content:
            return {
                "messages": [AIMessage(
                    content="",
                    tool_calls=[{
                        "name": "get_news",
                        "args": {"query": "AI新闻"},
                        "id": "call_001"
                    }]
                )]
            }
        # 分支2: 溢出直接结束
        if "overflow" in content:
            return {
                "messages": [AIMessage(content="上下文窗口溢出，已终止执行。")]
            }
        return None

    @hook_configure(can_jump_to="model")
    def after_model(self, state, runtime):
        """二次重试逻辑"""
        messages = state["messages"]
        user_msgs = [m for m in messages if isinstance(m, HumanMessage)]
        if not user_msgs:
            return None
        last_user = user_msgs[-1]
        if "retry_model" not in last_user.content:
            return None
        ai_msgs = [m for m in messages if isinstance(m, AIMessage)]
        if ai_msgs and "二次回答" in ai_msgs[-1].content:
            return None
        return {
            "messages": [SystemMessage(
                content="你必须以'二次回答'开头，并且只用一句话回答。"
            )]
        }

# 使用
agent = create_agent(
    model=model,
    tools=[get_news],
    middleware=[MyMiddleware()]
)