from typing import Literal, Annotated
from langchain_core.tools import tool
from langgraph.types import Command
from langgraph.graph import StateGraph, MessagesState
from langgraph.prebuilt import InjectedState

# 創建handoff工具的函數
def create_handoff_tool(*, agent_name: str, description: str = None):
    name = f"transfer_to_{agent_name}"
    description = description orf"Ask {agent_name} for help."
    
    @tool(name, description=description)
    def handoff_tool(
        task_description: Annotated[
            str, 
            "Description of what the next agent should do"
        ],
        state: Annotated[MessagesState, InjectedState],
    ) -> Command:
        # 創建任務描述消息
        task_message = {"role": "user", "content": task_description}
        # 準備下一個智能體的輸入
        agent_input = {**state, "messages": [task_message]}
        
        return Command(
            goto=agent_name,  # 指定下一個智能體
            update=agent_input,  # 傳遞的狀態更新
            graph=Command.PARENT  # 在父圖中導航
        )
    
    return handoff_tool

def agent_node(state: MessagesState) -> Command[Literal["agent_a", "agent_b", "END"]]:
    # 智能體邏輯處理
    messages = state["messages"]
    
    # 基于某種條件決定路由
    if need_expert_help(messages):
        return Command(
            goto="agent_b",  # 路由到agent_b
            update={"messages": messages + [new_message]}
        )
    elif task_completed(messages):
        return Command(
            goto="END",
            update={"messages": messages + [final_message]}
        )
    else:
        # 繼續當前智能體的處理
        return Command(
            goto="agent_a",  # 路由回自己
            update={"messages": messages + [processing_message]}
        )

from langgraph_supervisor import create_supervisor
from langchain.chat_models import init_chat_model

# 初始化模型
supervisor_model = init_chat_model("openai:gpt-4")

# 創建專業智能體
research_agent = create_react_agent(
    model="openai:gpt-4",
    tools=[web_search_tool, document_reader_tool],
    prompt="You are a research expert. Focus on finding accurate information.",
    name="research_expert"
)

math_agent = create_react_agent(
    model="openai:gpt-4",
    tools=[calculator_tool, equation_solver_tool],
    prompt="You are a mathematics expert. Solve complex calculations.",
    name="math_expert"
)

# 創建主管系統
supervisor_app = create_supervisor(
    agents=[research_agent, math_agent],
    model=supervisor_model,
    prompt=(
        "You are a supervisor managing two agents:\n"
        "- a research agent for information gathering\n"
        "- a math agent for calculations\n"
        "Assign work to one agent at a time based on the task."
    ),
    add_handoff_back_messages=True,
    output_mode="full_history"
).compile()

from typing import Annotated
from langchain_core.tools import tool, BaseTool, InjectedToolCallId
from langchain_core.messages import ToolMessage
from langgraph.types import Command
from langgraph.prebuilt import InjectedState

def create_custom_handoff_tool(
    *, 
    agent_name: str, 
    name: str = None, 
    description: str = None
) -> BaseTool:
    
    @tool(name, description=description)
    def handoff_to_agent(
        # LLM填充的任務描述
        task_description: Annotated[
            str, 
            "Detailed description of the task including context"
        ],
        # 可以添加額外的參數
        priority: Annotated[
            str,
            "Priority level: high, medium, low"
        ],
        # 注入的狀態
        state: Annotated[dict, InjectedState],
        tool_call_id: Annotated[str, InjectedToolCallId],
    ):
        # 創建工具消息
        tool_message = ToolMessage(
            content=f"Successfully transferred to {agent_name} with priority {priority}",
            name=name,
            tool_call_id=tool_call_id,
        )
        
        # 更新狀態并路由
        return Command(
            goto=agent_name,
            update={
                "messages": state["messages"] + [tool_message],
                "current_task": task_description,
                "priority": priority
            },
            graph=Command.PARENT
        )
    
    return handoff_to_agent

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.store.memory import InMemoryStore

# 短期記憶（對話狀態）
checkpointer = InMemorySaver()
# 長期記憶（知識存儲）
store = InMemoryStore()

# 編譯時添加持久化支持
app = supervisor_app.compile(
    checkpointer=checkpointer,
    store=store
)

# 創建主管時指定輸出模式
supervisor = create_supervisor(
    agents=[agent1, agent2],
    model=model,
    output_mode="last_message",  # 只返回最后的消息
    # output_mode="full_history"  # 返回完整歷史
)

from langgraph_supervisor.handoff import create_forward_message_tool

# 創建轉發工具
forwarding_tool = create_forward_message_tool("supervisor")

# 在創建主管時添加
workflow = create_supervisor(
    [research_agent, math_agent],
    model=model,
    tools=[forwarding_tool]  # 添加轉發工具
)

class MultiAgentState(TypedDict):
    messages: Annotated[list, add_messages]
    current_agent: str
    task_queue: list
    results: dict

def create_complex_workflow():
    workflow = StateGraph(MultiAgentState)
    
    # 添加智能體節點
    workflow.add_node("supervisor", supervisor_agent)
    workflow.add_node("researcher", research_agent)
    workflow.add_node("analyst", data_analyst_agent)
    workflow.add_node("writer", content_writer_agent)
    
    # 使用Command進行動態路由
    def supervisor_agent(state: MultiAgentState) -> Command:
        task = state["task_queue"][0] if state["task_queue"] elseNone
        
        ifnot task:
            return Command(goto="END", update=state)
        
        # 基于任務類型分配給不同智能體
        if task["type"] == "research":
            return Command(
                goto="researcher",
                update={**state, "current_agent": "researcher"}
            )
        elif task["type"] == "analysis":
            return Command(
                goto="analyst",
                update={**state, "current_agent": "analyst"}
            )
        # ... 更多路由邏輯
    
    return workflow.compile()

def selective_handoff(state: MessagesState) -> Command:
    # 只傳遞最近的N條消息
    recent_messages = state["messages"][-10:]
    
    # 創建摘要
    summary = create_summary(state["messages"][:-10])
    
    return Command(
        goto="next_agent",
        update={"messages": [summary] + recent_messages}
    )

def robust_agent(state: MessagesState) -> Command:
    try:
        # 智能體邏輯
        result = process_task(state)
        return Command(goto="success_handler", update={"result": result})
    except Exception as e:
        # 錯誤恢復
        return Command(
            goto="error_handler",
            update={"error": str(e), "fallback_agent": "supervisor"}
        )

from langgraph_supervisor import create_supervisor
from langgraph.prebuilt import create_react_agent
import os

# 設置API密鑰
os.environ["OPENAI_API_KEY"] = "your-api-key"

# 1. 創建專業智能體
# FAQ智能體
faq_agent = create_react_agent(
    model="openai:gpt-3.5-turbo",
    tools=[search_faq_tool, get_product_info_tool],
    prompt="You handle frequently asked questions about products and services.",
    name="faq_specialist"
)

# 技術支持智能體
tech_agent = create_react_agent(
    model="openai:gpt-4",
    tools=[diagnose_issue_tool, check_system_status_tool],
    prompt="You are a technical support specialist. Help users solve technical problems.",
    name="tech_specialist"
)

# 訂單處理智能體
order_agent = create_react_agent(
    model="openai:gpt-3.5-turbo",
    tools=[check_order_status_tool, process_return_tool],
    prompt="You handle order-related inquiries and returns.",
    name="order_specialist"
)

# 2. 創建主管系統
customer_service_system = create_supervisor(
    agents=[faq_agent, tech_agent, order_agent],
    model="openai:gpt-4",
    prompt="""
    You are a customer service supervisor managing three specialists:
    - FAQ specialist: handles general questions
    - Tech specialist: handles technical issues
    - Order specialist: handles order and shipping issues
    
    Analyze the customer's query and delegate to the appropriate specialist.
    If the query spans multiple areas, handle them sequentially.
    """,
    add_handoff_back_messages=True,
    output_mode="last_message"# 返回最終響應
)

# 3. 編譯并運行
app = customer_service_system.compile(
    checkpointer=InMemorySaver()  # 添加對話記憶
)

# 4. 處理客戶查詢
asyncdef handle_customer_query(query: str, thread_id: str):
    config = {"configurable": {"thread_id": thread_id}}
    
    result = await app.ainvoke(
        {"messages": [{"role": "user", "content": query}]},
        config=config
    )
    
    return result["messages"][-1]["content"]

# 使用示例
response = await handle_customer_query(
    "My laptop won't turn on and I want to check if it's still under warranty",
    thread_id="customer_123"
)