import json
import operator
from typing import Annotated, TypedDict, Union, List

from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langgraph.graph import StateGraph, END, START
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import MemorySaver


@tool
def start_stream(stream_id: str, protocol: str):
    """Start a stream with given stream_id and protocol (rtmp/hls only)."""
    if protocol not in ['rtmp', 'hls']:
        return json.dumps({"error": "Unsupported protocol"})
    return json.dumps({"stream_id": stream_id, "protocol": protocol, "status": "starting"})

@tool
def stop_stream(stream_id: str):
    """Stop a running stream."""
    return json.dumps({"msg": f"stream_{stream_id} has been stopped."})

@tool
def get_stream_status(stream_id: str):
    """Get technical status (bitrate, fps) of a stream."""
    if stream_id == "news_channel":
        return json.dumps({"bitrate": 500, "status": "unstable", "fps": 10})
    elif stream_id == "sports_channel":
        return json.dumps({"bitrate": 4000, "status": "live", "fps": 60})
    else:
        return json.dumps({"error": "offline"})

tools = [start_stream, stop_stream, get_stream_status]


# Annotated[list, operator.add]：新的消息会 "Append" 到列表，而不是覆盖
class AgentState(TypedDict):
    messages: Annotated[List[BaseMessage], operator.add]


llm = ChatOpenAI(
    model="gpt-4o-mini",
    api_key=,
    base_url=
)
llm_with_tools = llm.bind_tools(tools)

def agent_node(state: AgentState):
    """思考节点：调用 LLM 决定下一步"""
    messages = state["messages"]
    response = llm_with_tools.invoke(messages)
    # 返回新的消息（LangGraph 会自动将其 append 到 state）
    return {"messages": [response]}


tool_node = ToolNode(tools)

# 定义图
workflow = StateGraph(AgentState)

# 添加节点
workflow.add_node("agent", agent_node)
workflow.add_node("tools", tool_node)

# 定义边 (Edges)
workflow.add_edge(START, "agent") # 入口 -> 思考

# 定义条件边 (Conditional Edges)
# 逻辑：如果 agent 输出包含 tool_calls，跳转到 "tools" 节点，否则结束
def should_continue(state: AgentState):
    last_message = state["messages"][-1]
    if last_message.tool_calls:
        return "tools"
    return END

workflow.add_conditional_edges(
    "agent",
    should_continue,
    ["tools", END]
)

# 工具执行完后，必须跳回 agent 继续思考
workflow.add_edge("tools", "agent") 

# memory=checkpointer 开启了持久化记忆
checkpointer = MemorySaver()
app = workflow.compile(checkpointer=checkpointer)

# 模拟运行
if __name__ == "__main__":
    # 配置线程 ID，模拟不同用户的会话
    config = {"configurable": {"thread_id": "user_session_001"}}
    
    print("--- 任务：复杂流控推理 ---")
    user_input = "检查一下 news_channel 的状态，如果它不稳定或者码率低于 1000，就把它重启一下。"
    
    # 这一步包含了整个 While 循环 + 状态管理
    # Stream 模式允许我们看到中间步骤
    for event in app.stream(
        {"messages": [HumanMessage(content=user_input)]}, 
        config=config
    ):
        for key, value in event.items():
            print(f"\n[Node: {key}]")
            # 打印最新的消息内容
            last_msg = value["messages"][-1]
            if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
                print(f"  -> 计划调用: {last_msg.tool_calls[0]['name']}")
            elif hasattr(last_msg, "content"):
                 print(f"  -> 内容: {last_msg.content}")

    print("\n\n====== 最终结果 ======")
    # 由于有 checkpointer，我们可以随时获取当前状态
    final_state = app.get_state(config)
    print(final_state.values["messages"][-1].content)

# ... (前面的代码保持不变) ...

# 1. 修改编译步骤：增加 interrupt_before
# 这意味着在进入 "tools" 节点前，程序会暂停
app = workflow.compile(
    checkpointer=checkpointer,
    interrupt_before=["tools"] # <--- 关键修改
)

if __name__ == "__main__":
    config = {"configurable": {"thread_id": "user_session_hitl_01"}}
    
    print("--- 阶段1：触发任务 ---")
    user_input = "news_channel 好像卡了，帮我重启一下。"
    
    # 第一次运行：Agent 会思考并决定调用工具，但会在执行前暂停
    for event in app.stream({"messages": [HumanMessage(content=user_input)]}, config=config):
        for key, value in event.items():
            print(f"[Node: {key}] 正在处理...")

    # 此时，程序暂停了。我们可以检查 Agent 到底想干什么
    snapshot = app.get_state(config)
    next_step = snapshot.next
    
    if "tools" in next_step:
        last_msg = snapshot.values["messages"][-1]
        tool_call = last_msg.tool_calls[0]
        print(f"\n[系统暂停] Agent 申请执行: {tool_call['name']} 参数: {tool_call['args']}")
        
        # 模拟人工审批
        approval = input(">>> 是否批准执行？(y/n): ")
        
        if approval.lower() == 'y':
            print("--- 阶段2：批准并继续 ---")
            # 传入 None 表示继续执行当前状态，不添加新消息
            for event in app.stream(None, config=config):
                for key, value in event.items():
                     print(f"[Node: {key}] {value}")
        else:
            print("--- 阶段2：拒绝并告知 Agent ---")
            # 我们可以直接注入一条 ToolMessage 表示拒绝，或者注入一条 HumanMessage 告诉它不行
            app.update_state(
                config,
                {"messages": [HumanMessage(content="管理员拒绝了重启操作，请检查是否有其他非破坏性方案。")]}
            )
            # 继续运行，Agent 会收到拒绝信息并重新思考
            for event in app.stream(None, config=config):
                 print(f"[Node: {key}] {value}")

from langchain_core.messages import RemoveMessage, SystemMessage, HumanMessage

def summarize_node(state: AgentState):
    summary_prompt = "请总结以上对话的关键信息..."
    messages = state["messages"]
    
    # 调用 LLM 生成总结
    summary = llm.invoke(messages + [HumanMessage(content=summary_prompt)])
    
    # 找出除了最新的 2 条消息以外的所有旧消息 ID
    # 我们保留最后 2 条是为了上下文连贯性
    messages_to_delete = [m.id for m in messages[:-2]]
    
    # 构造删除操作
    delete_operations = [RemoveMessage(id=mid) for mid in messages_to_delete]
    
    # 将总结作为新的 SystemMessage 插入（或者作为一段上下文）
    new_summary_msg = SystemMessage(content=f"Background Summary: {summary.content}")
    
    # 返回：删除旧的 + 插入新的总结
    return {"messages": [new_summary_msg] + delete_operations}

# 需要在 workflow 中添加条件边
def should_summarize(state: AgentState):
    if len(state["messages"]) > 10:
        return "summarize"
    return END # 或者继续其他逻辑

from langchain_core.messages import trim_messages

# 定义裁剪器
trimmer = trim_messages(
    max_tokens=2000,
    strategy="last",
    token_counter=llm, # 使用模型的 tokenizer 计数
    include_system=True, # 总是保留 System Prompt
    allow_partial=False,
    start_on="human" # 确保截断后的第一条是人类消息（避免从中间的 tool 切断）
)

def agent_node(state: AgentState):
    # 在调用 LLM 前先过一遍 trimmer
    trimmed_messages = trimmer.invoke(state["messages"])
    response = llm_with_tools.invoke(trimmed_messages)
    return {"messages": [response]}