from typing import TypedDict
from loguru import logger
import time
from langgraph.types import interrupt, Command


# https://xbsheng.github.io/atguigu-note/langgraph/%E8%AF%BE%E4%BB%B6/05-LangGraph%E9%AB%98%E7%BA%A7%E7%89%B9%E6%80%A7


class OverallState(TypedDict):
    initial_input: str
    node_a1: str
    node_a2: str
    node_b: str


def node_a1(state: OverallState) -> OverallState:
    time.sleep(1)
    logger.info("node_a1")
    return OverallState(node_a1="node_a1")


def node_a2(state: OverallState) -> OverallState:
    time.sleep(3)
    logger.info("node_a2")

    return OverallState(node_a2="node_a2")


def node_b(state: OverallState) -> OverallState:
    time.sleep(1)
    res = interrupt("hello")
    logger.info("node_b")
    return OverallState(node_b=f"node_b,res={res}")


from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import InMemorySaver

builder = StateGraph(state_schema=OverallState)
builder.add_node("node_a1", node_a1)
builder.add_node("node_a2", node_a2)
builder.add_node("node_b", node_b)

builder.add_edge(START, "node_a1")
builder.add_edge(START, "node_a2")
builder.add_edge(["node_a1", "node_a2"], "node_b")
builder.add_edge("node_b", END)

checkpoint_saver = InMemorySaver()

graph = builder.compile(checkpointer=checkpoint_saver)

from IPython.display import display

display(graph)

from rich import print as rp

config = {
    "configurable": {
        "thread_id": "12345"
    }
}

for chunk in graph.stream(
        {"initial_input": "init"},
        stream_mode=["debug"],
        config=config
):
    print("=" * 50)
    rp(chunk)

# print("=" * 50)
# rp(list(graph.get_state_history(config=config)))
print("**" * 50)

for chunk in graph.stream(
        Command(resume="interrupt_node_b_res"),
        stream_mode=["debug"],
        config=config
):
    print("=" * 50)
    rp(chunk)
# print("=" * 50)
# rp(list(graph.get_state_history(config=config)))