from typing import Annotated, Literal, Optional

from typing_extensions import TypedDict

from langgraph.graph.message import AnyMessage, add_messages


def update_dialog_stack(left: list[str], right: Optional[str]) -> list[str]:
    """推入或彈出狀態。"""
    if right is None:
        return left
    if right == "pop":
        return left[:-1]
    return left + [right]


class State(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    user_info: str
    dialog_state: Annotated[
        list[
            Literal[
                "assistant",
                "update_flight",
                "book_car_rental",
                "book_hotel",
                "book_excursion",
            ]
        ],
        update_dialog_stack,
    ]

from langchain_anthropic import ChatAnthropic
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.runnables import Runnable, RunnableConfig


class Assistant:
    def __init__(self, runnable: Runnable):
        self.runnable = runnable

    def __call__(self, state: State, config: RunnableConfig):
        while True:
            result = self.runnable.invoke(state)

            if not result.tool_calls and (
                not result.content
                or isinstance(result.content, list)
                and not result.content[0].get("text")
            ):
                messages = state["messages"] + [("user", "用真實的輸出回應。")]
                state = {**state, "messages": messages}
                messages = state["messages"] + [("user", "用真實的輸出回應。")]
                state = {**state, "messages": messages}
            else:
                break
        return {"messages": result}


class CompleteOrEscalate(BaseModel):
    """一個工具，標記當前任務為已完成和/或將對話控制權升級到主助手，
    主助手可以根據用戶的需求重新路由對話。"""

    cancel: bool = True
    reason: str

    class Config:
        schema_extra = {
            "example": {
                "cancel": True,
                "reason": "用戶改變了他們對當前任務的想法。",
            },
            "example 2": {
                "cancel": True,
                "reason": "我已經完全完成了任務。",
            },
            "example 3": {
                "cancel": False,
                "reason": "我需要搜索用戶的電子郵件或日歷以獲取更多信息。",
            },
        }

# 航班預訂助手

flight_booking_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "您是專門處理航班更新的助手。"
            "每當用戶需要幫助更新他們的預訂時，主要助手就會委派工作給您。"
            "與客戶確認更新的航班詳情，并告知他們任何額外費用。"
            "搜索時，要堅持不懈。如果第一次搜索沒有結果，就擴大您的查詢范圍。"
            "如果您需要更多信息或客戶改變了主意，將任務升級回主助手。"
            "記住，只有在相關工具成功使用后，預訂才算完成。"
            "\n\n當前用戶航班信息：\n<Flights>\n{user_info}\n</Flights>"
            "\n當前時間：{time}。"
            "\n\n如果用戶需要幫助，而且您的工具都不適合，那么"
            '"CompleteOrEscalate" 對話到主機助手。不要浪費用戶的時間。不要編造無效的工具或功能。',
        ),
        ("placeholder", "{messages}"),
    ]
).partial(time=datetime.now())

update_flight_safe_tools = [search_flights]
update_flight_sensitive_tools = [update_ticket_to_new_flight, cancel_ticket]
update_flight_tools = update_flight_safe_tools + update_flight_sensitive_tools
update_flight_runnable = flight_booking_prompt | llm.bind_tools(
    update_flight_tools + [CompleteOrEscalate]
)

# 汽車租賃助手
book_car_rental_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "您是專門處理汽車租賃預訂的助手。"
            "每當用戶需要幫助預訂汽車租賃時，主要助手就會委派工作給您。"
            "根據用戶的偏好搜索可用的汽車租賃，并與客戶確認預訂詳情。"
            "搜索時，要堅持不懈。如果第一次搜索沒有結果，就擴大您的查詢范圍。"
            "如果您需要更多信息或客戶改變了主意，將任務升級回主助手。"
            "記住，只有在相關工具成功使用后，預訂才算完成。"
            "\n當前時間：{time}。"
            "\n\n如果用戶需要幫助，而且您的工具都不適合，那么 "
            '"CompleteOrEscalate" 對話到主機助手。不要浪費用戶的時間。不要編造無效的工具或功能。'
            "\n\n一些您應該 CompleteOrEscalate 的示例："
            " - '今年這個時候的天氣怎么樣？'"
            " - '有哪些航班可用？'"
            " - '沒關系，我想我會單獨預訂'"
            " - '哦等等，我還沒有預訂我的航班，我會先做這件事'"
            " - '汽車租賃預訂確認'",
        ),
        ("placeholder", "{messages}"),
    ]
).partial(time=datetime.now())

book_car_rental_safe_tools = [search_car_rentals]
book_car_rental_sensitive_tools = [
    book_car_rental,
    update_car_rental,
    cancel_car_rental,
]
book_car_rental_tools = book_car_rental_safe_tools + book_car_rental_sensitive_tools
book_car_rental_runnable = book_car_rental_prompt | llm.bind_tools(
    book_car_rental_tools + [CompleteOrEscalate]
)

# 酒店預訂助手
book_hotel_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "您是專門處理酒店預訂的助手。"
            "每當用戶需要幫助預訂酒店時，主要助手就會委派工作給您。"
            "根據用戶的偏好搜索可用的酒店，并與客戶確認預訂詳情。"
            "搜索時，要堅持不懈。如果第一次搜索沒有結果，就擴大您的查詢范圍。"
            "如果您需要更多信息或客戶改變了主意，將任務升級回主助手。"
            "記住，只有在相關工具成功使用后，預訂才算完成。"
            "\n當前時間：{time}。"
            '\n\n如果用戶需要幫助，而且您的工具都不適合，那么 "CompleteOrEscalate" 對話到主機助手。'
            "不要浪費用戶的時間。不要編造無效的工具或功能。"
            "\n\n一些您應該 CompleteOrEscalate 的示例："
            " - '今年這個時候的天氣怎么樣？'"
            " - '沒關系，我想我會單獨預訂'"
            " - '我需要弄清楚我在那里的時候的交通'"
            " - '哦等等，我還沒有預訂我的航班，我會先做這件事'"
            " - '酒店預訂確認'",
        ),
        ("placeholder", "{messages}"),
    ]
).partial(time=datetime.now())

book_hotel_safe_tools = [search_hotels]
book_hotel_sensitive_tools = [book_hotel, update_hotel, cancel_hotel]
book_hotel_tools = book_hotel_safe_tools + book_hotel_sensitive_tools
book_hotel_runnable = book_hotel_prompt | llm.bind_tools(
    book_hotel_tools + [CompleteOrEscalate]
)

# 旅行助手

book_excursion_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "您是專門處理旅行建議的助手。"
            "每當用戶需要幫助預訂推薦的旅行時，主要助手就會委派工作給您。"
            "根據用戶的偏好搜索可用的旅行建議，并與客戶確認預訂詳情。"
            "如果您需要更多信息或客戶改變了主意，將任務升級回主助手。"
            "搜索時，要堅持不懈。如果第一次搜索沒有結果，就擴大您的查詢范圍。"
            "記住，只有在相關工具成功使用后，預訂才算完成。"
            "\n當前時間：{time}。"
            '\n\n如果用戶需要幫助，而且您的工具都不適合，那么 "CompleteOrEscalate" 對話到主機助手。不要浪費用戶的時間。不要編造無效的工具或功能。'
            "\n\n一些您應該 CompleteOrEscalate 的示例："
            " - '沒關系，我想我會單獨預訂'"
            " - '我需要在那里的時候弄清楚交通'"
            " - '哦等等，我還沒有預訂我的航班，我會先做這件事'"
            " - '旅行預訂確認！'",
        ),
        ("placeholder", "{messages}"),
    ]
).partial(time=datetime.now())

book_excursion_safe_tools = [search_trip_recommendations]
book_excursion_sensitive_tools = [book_excursion, update_excursion, cancel_excursion]
book_excursion_tools = book_excursion_safe_tools + book_excursion_sensitive_tools
book_excursion_runnable = book_excursion_prompt | llm.bind_tools(
    book_excursion_tools + [CompleteOrEscalate]
)

# 主要助手
class ToFlightBookingAssistant(BaseModel):
    """將工作轉交給專門助手來處理航班更新和取消。"""

    request: str = Field(
        descriptinotallow="更新航班助手在繼續操作之前需要澄清的任何必要的后續問題。"
    )


class ToBookCarRental(BaseModel):
    """將工作轉交給專門助手來處理汽車租賃預訂。"""

    location: str = Field(
        descriptinotallow="用戶想要租車的地點。"
    )
    start_date: str = Field(descriptinotallow="汽車租賃的開始日期。")
    end_date: str = Field(descriptinotallow="汽車租賃的結束日期。")
    request: str = Field(
        descriptinotallow="用戶關于汽車租賃的任何額外信息或請求。"
    )

    class Config:
        schema_extra = {
            "example": {
                "location": "巴塞爾",
                "start_date": "2023-07-01",
                "end_date": "2023-07-05",
                "request": "我需要一輛自動變速箱的緊湊型汽車。",
            }
        }


class ToHotelBookingAssistant(BaseModel):
    """將工作轉交給專門助手來處理酒店預訂。"""

    location: str = Field(
        descriptinotallow="用戶想要預訂酒店的地點。"
    )
    checkin_date: str = Field(descriptinotallow="酒店的入住日期。")
    checkout_date: str = Field(descriptinotallow="酒店的退房日期。")
    request: str = Field(
        descriptinotallow="用戶關于酒店預訂的任何額外信息或請求。"
    )

    class Config:
        schema_extra = {
            "example": {
                "location": "蘇黎世",
                "checkin_date": "2023-08-15",
                "checkout_date": "2023-08-20",
                "request": "我更喜歡靠近市中心的酒店，有風景的房間。",
            }
        }


class ToBookExcursion(BaseModel):
    """將工作轉交給專門助手來處理旅行推薦和其他旅行預訂。"""

    location: str = Field(
        descriptinotallow="用戶想要預訂推薦旅行的地點。"
    )
    request: str = Field(
        descriptinotallow="用戶關于旅行建議的任何額外信息或請求。"
    )

    class Config:
        schema_extra = {
            "example": {
                "location": "盧塞恩",
                "request": "用戶對戶外活動和風景感興趣。",
            }
        }


# 最高級助手執行一般問答，并將專業任務委派給其他助手。
# 任務委派是一種簡單的語義路由/簡單的意圖檢測
# llm = ChatAnthropic(model="claude-3-haiku-20240307")
llm = ChatAnthropic(model="claude-3-sonnet-20240229", temperature=1)

primary_assistant_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "您是瑞士航空的樂于助人的客服助手。"
            "您的主要角色是搜索航班信息和公司政策，以回答客戶查詢。"
            "如果客戶請求更新或取消預訂、租車、預訂酒店或獲取旅行建議，"
            "通過調用相應的工具，將任務委派給適當的專業助手。您自己無法進行這些類型的更改。"
            "只有專業助手才有權為用戶執行此操作。"
            "用戶不知道不同的專業助手，所以不要提及他們；只需通過功能調用來靜靜地委派。"
            "向客戶提供詳細信息，并在得出信息不可用的結論之前，始終再次檢查數據庫。"
            "當搜索時，要堅持不懈。如果第一次搜索沒有結果，就擴大您的查詢范圍。"
            "\n\n當前用戶航班信息：\n<Flights>\n{user_info}\n</Flights>"
            "\n當前時間：{time}。",
        ),
        ("placeholder", "{messages}"),
    ]
).partial(time=datetime.now())
primary_assistant_tools = [
    TavilySearchResults(max_results=1),
    search_flights,
    lookup_policy,
]
assistant_runnable = primary_assistant_prompt | llm.bind_tools(
    primary_assistant_tools
    + [
        ToFlightBookingAssistant,
        ToBookCarRental,
        ToHotelBookingAssistant,
        ToBookExcursion,
    ]
)

from typing import Callable

from langchain_core.messages import ToolMessage


def create_entry_node(assistant_name: str, new_dialog_state: str) -> Callable:
    def entry_node(state: State) -> dict:
        tool_call_id = state["messages"][-1].tool_calls[0]["id"]
        return {
            "messages": [
                ToolMessage(
                    cnotallow=f"現在助手是 {assistant_name}。回想一下主機助手和用戶之間的上述對話。"
                    f" 用戶的意圖尚未得到滿足。使用提供的工具來幫助用戶。記住，你是 {assistant_name}，"
                    " 預訂、更新、其他操作或其他動作只有在您成功調用適當的工具后才完成。"
                    " 如果用戶改變主意或需要其他任務的幫助，請調用 CompleteOrEscalate 函數，讓主機助手接管控制權。"
                    " 不要提及你是誰 - 只作為助手的代理行事。",
                    tool_call_id=tool_call_id,
                )
            ],
            "dialog_state": new_dialog_state,
        }

    return entry_node

from typing import Literal

from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import END, StateGraph
from langgraph.prebuilt import tools_condition

builder = StateGraph(State)


def user_info(state: State):
    return {"user_info": fetch_user_flight_information.invoke({})}


builder.add_node("fetch_user_info", user_info)
builder.set_entry_point("fetch_user_info")

# 航班預訂助手
# Flight booking assistant
builder.add_node(
    "enter_update_flight",
    create_entry_node("Flight Updates & Booking Assistant", "update_flight"),
)
builder.add_node("update_flight", Assistant(update_flight_runnable))
builder.add_edge("enter_update_flight", "update_flight")
builder.add_node(
    "update_flight_sensitive_tools",
    create_tool_node_with_fallback(update_flight_sensitive_tools),
)
builder.add_node(
    "update_flight_safe_tools",
    create_tool_node_with_fallback(update_flight_safe_tools),
)


def route_update_flight(
    state: State,
) -> Literal[
    "update_flight_sensitive_tools",
    "update_flight_safe_tools",
    "leave_skill",
    "__end__",
]:
    route = tools_condition(state)
    if route == END:
        return END
    tool_calls = state["messages"][-1].tool_calls
    did_cancel = any(tc["name"] == CompleteOrEscalate.__name__ for tc in tool_calls)
    if did_cancel:
        return "leave_skill"
    safe_toolnames = [t.name for t in update_flight_safe_tools]
    if all(tc["name"] in safe_toolnames for tc in tool_calls):
        return "update_flight_safe_tools"
    return "update_flight_sensitive_tools"


builder.add_edge("update_flight_sensitive_tools", "update_flight")
builder.add_edge("update_flight_safe_tools", "update_flight")
builder.add_conditional_edges("update_flight", route_update_flight)


# This node will be shared for exiting all specialized assistants
def pop_dialog_state(state: State) -> dict:
    """Pop the dialog stack and return to the main assistant.

    This lets the full graph explicitly track the dialog flow and delegate control
    to specific sub-graphs.
    """
    messages = []
    if state["messages"][-1].tool_calls:
        # Note: Doesn't currently handle the edge case where the llm performs parallel tool calls
        messages.append(
            ToolMessage(
                cnotallow="Resuming dialog with the host assistant. Please reflect on the past conversation and assist the user as needed.",
                tool_call_id=state["messages"][-1].tool_calls[0]["id"],
            )
        )
    return {
        "dialog_state": "pop",
        "messages": messages,
    }


builder.add_node("leave_skill", pop_dialog_state)
builder.add_edge("leave_skill", "primary_assistant")

# 租車助手

# Car rental assistant

builder.add_node(
    "enter_book_car_rental",
    create_entry_node("Car Rental Assistant", "book_car_rental"),
)
builder.add_node("book_car_rental", Assistant(book_car_rental_runnable))
builder.add_edge("enter_book_car_rental", "book_car_rental")
builder.add_node(
    "book_car_rental_safe_tools",
    create_tool_node_with_fallback(book_car_rental_safe_tools),
)
builder.add_node(
    "book_car_rental_sensitive_tools",
    create_tool_node_with_fallback(book_car_rental_sensitive_tools),
)


def route_book_car_rental(
    state: State,
) -> Literal[
    "book_car_rental_safe_tools",
    "book_car_rental_sensitive_tools",
    "leave_skill",
    "__end__",
]:
    route = tools_condition(state)
    if route == END:
        return END
    tool_calls = state["messages"][-1].tool_calls
    did_cancel = any(tc["name"] == CompleteOrEscalate.__name__ for tc in tool_calls)
    if did_cancel:
        return "leave_skill"
    safe_toolnames = [t.name for t in book_car_rental_safe_tools]
    if all(tc["name"] in safe_toolnames for tc in tool_calls):
        return "book_car_rental_safe_tools"
    return "book_car_rental_sensitive_tools"


builder.add_edge("book_car_rental_sensitive_tools", "book_car_rental")
builder.add_edge("book_car_rental_safe_tools", "book_car_rental")
builder.add_conditional_edges("book_car_rental", route_book_car_rental)

# 酒店預訂助手
# Hotel booking assistant
builder.add_node(
    "enter_book_hotel", create_entry_node("Hotel Booking Assistant", "book_hotel")
)
builder.add_node("book_hotel", Assistant(book_hotel_runnable))
builder.add_edge("enter_book_hotel", "book_hotel")
builder.add_node(
    "book_hotel_safe_tools",
    create_tool_node_with_fallback(book_hotel_safe_tools),
)
builder.add_node(
    "book_hotel_sensitive_tools",
    create_tool_node_with_fallback(book_hotel_sensitive_tools),
)


def route_book_hotel(
    state: State,
) -> Literal[
    "leave_skill", "book_hotel_safe_tools", "book_hotel_sensitive_tools", "__end__"
]:
    route = tools_condition(state)
    if route == END:
        return END
    tool_calls = state["messages"][-1].tool_calls
    did_cancel = any(tc["name"] == CompleteOrEscalate.__name__ for tc in tool_calls)
    if did_cancel:
        return "leave_skill"
    tool_names = [t.name for t in book_hotel_safe_tools]
    if all(tc["name"] in tool_names for tc in tool_calls):
        return "book_hotel_safe_tools"
    return "book_hotel_sensitive_tools"


builder.add_edge("book_hotel_sensitive_tools", "book_hotel")
builder.add_edge("book_hotel_safe_tools", "book_hotel")
builder.add_conditional_edges("book_hotel", route_book_hotel)

# 旅行預訂助手
# Excursion assistant
builder.add_node(
    "enter_book_excursion",
    create_entry_node("Trip Recommendation Assistant", "book_excursion"),
)
builder.add_node("book_excursion", Assistant(book_excursion_runnable))
builder.add_edge("enter_book_excursion", "book_excursion")
builder.add_node(
    "book_excursion_safe_tools",
    create_tool_node_with_fallback(book_excursion_safe_tools),
)
builder.add_node(
    "book_excursion_sensitive_tools",
    create_tool_node_with_fallback(book_excursion_sensitive_tools),
)


def route_book_excursion(
    state: State,
) -> Literal[
    "book_excursion_safe_tools",
    "book_excursion_sensitive_tools",
    "leave_skill",
    "__end__",
]:
    route = tools_condition(state)
    if route == END:
        return END
    tool_calls = state["messages"][-1].tool_calls
    did_cancel = any(tc["name"] == CompleteOrEscalate.__name__ for tc in tool_calls)
    if did_cancel:
        return "leave_skill"
    tool_names = [t.name for t in book_excursion_safe_tools]
    if all(tc["name"] in tool_names for tc in tool_calls):
        return "book_excursion_safe_tools"
    return "book_excursion_sensitive_tools"


builder.add_edge("book_excursion_sensitive_tools", "book_excursion")
builder.add_edge("book_excursion_safe_tools", "book_excursion")
builder.add_conditional_edges("book_excursion", route_book_excursion)

# Primary assistant
builder.add_node("primary_assistant", Assistant(assistant_runnable))
builder.add_node(
    "primary_assistant_tools", create_tool_node_with_fallback(primary_assistant_tools)
)


def route_primary_assistant(
    state: State,
) -> Literal[
    "primary_assistant_tools",
    "enter_update_flight",
    "enter_book_hotel",
    "enter_book_excursion",
    "__end__",
]:
    route = tools_condition(state)
    if route == END:
        return END
    tool_calls = state["messages"][-1].tool_calls
    if tool_calls:
        if tool_calls[0]["name"] == ToFlightBookingAssistant.__name__:
            return "enter_update_flight"
        elif tool_calls[0]["name"] == ToBookCarRental.__name__:
            return "enter_book_car_rental"
        elif tool_calls[0]["name"] == ToHotelBookingAssistant.__name__:
            return "enter_book_hotel"
        elif tool_calls[0]["name"] == ToBookExcursion.__name__:
            return "enter_book_excursion"
        return "primary_assistant_tools"
    raise ValueError("Invalid route")


# The assistant can route to one of the delegated assistants,
# directly use a tool, or directly respond to the user
builder.add_conditional_edges(
    "primary_assistant",
    route_primary_assistant,
    {
        "enter_update_flight": "enter_update_flight",
        "enter_book_car_rental": "enter_book_car_rental",
        "enter_book_hotel": "enter_book_hotel",
        "enter_book_excursion": "enter_book_excursion",
        "primary_assistant_tools": "primary_assistant_tools",
        END: END,
    },
)
builder.add_edge("primary_assistant_tools", "primary_assistant")


# Each delegated workflow can directly respond to the user
# When the user responds, we want to return to the currently active workflow
def route_to_workflow(
    state: State,
) -> Literal[
    "primary_assistant",
    "update_flight",
    "book_car_rental",
    "book_hotel",
    "book_excursion",
]:
    """If we are in a delegated state, route directly to the appropriate assistant."""
    dialog_state = state.get("dialog_state")
    if not dialog_state:
        return "primary_assistant"
    return dialog_state[-1]


builder.add_conditional_edges("fetch_user_info", route_to_workflow)

# Compile graph
memory = SqliteSaver.from_conn_string(":memory:")
part_4_graph = builder.compile(
    checkpointer=memory,
    # Let the user approve or deny the use of sensitive tools
    interrupt_before=[
        "update_flight_sensitive_tools",
        "book_car_rental_sensitive_tools",
        "book_hotel_sensitive_tools",
        "book_excursion_sensitive_tools",
    ],
)

import shutil
import uuid

# Update with the backup file so we can restart from the original place in each section
shutil.copy(backup_file, db)
thread_id = str(uuid.uuid4())

config = {
    "configurable": {
        # The passenger_id is used in our flight tools to
        # fetch the user's flight information
        "passenger_id": "3442 587242",
        # Checkpoints are accessed by thread_id
        "thread_id": thread_id,
    }
}

_printed = set()
# We can reuse the tutorial questions from part 1 to see how it does.
for question in tutorial_questions:
    events = part_4_graph.stream(
        {"messages": ("user", question)}, config, stream_mode="values"
    )
    for event in events:
        _print_event(event, _printed)
    snapshot = part_4_graph.get_state(config)
    while snapshot.next:
        # We have an interrupt! The agent is trying to use a tool, and the user can approve or deny it
        # Note: This code is all outside of your graph. Typically, you would stream the output to a UI.
        # Then, you would have the frontend trigger a new run via an API call when the user has provided input.
        user_input = input(
            "Do you approve of the above actions? Type 'y' to continue;"
            " otherwise, explain your requested changed.\n\n"
        )
        if user_input.strip() == "y":
            # Just continue
            result = part_4_graph.invoke(
                None,
                config,
            )
        else:
            # Satisfy the tool invocation by
            # providing instructions on the requested changes / change of mind
            result = part_4_graph.invoke(
                {
                    "messages": [
                        ToolMessage(
                            tool_call_id=event["messages"][-1].tool_calls[0]["id"],
                            cnotallow=f"API call denied by user. Reasoning: '{user_input}'. Continue assisting, accounting for the user's input.",
                        )
                    ]
                },
                config,
            )
        snapshot = part_4_graph.get_state(config)