"""Tests for subgraph persistence behavior (sync).

Covers three checkpointer settings for subgraph state:
- checkpointer=False: no persistence, even when parent has a checkpointer
- checkpointer=None (default): "stateless" — inherits parent checkpointer for
  interrupt support, but state resets each invocation. This is the common case
  when an agent is invoked from inside a tool used by another agent.
- checkpointer=True: "stateful" — state accumulates across invocations on the same thread id
"""

from uuid import uuid4

from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.base import BaseCheckpointSaver
from typing_extensions import TypedDict

from langgraph.graph import START, StateGraph
from langgraph.graph.message import MessagesState
from langgraph.types import Command, Interrupt, interrupt
from tests.any_str import AnyStr


class ParentState(TypedDict):
    result: str


# -- checkpointer=None (stateless) --


def test_stateless_interrupt_resume(
    sync_checkpointer: BaseCheckpointSaver,
) -> None:
    """Tests that a subgraph compiled with checkpointer=None (the default) can
    still support interrupt/resume when invoked from inside a parent graph that
    has a checkpointer. This is the "stateless" pattern — the subgraph inherits
    the parent's checkpointer just enough to pause and resume, but does not
    retain any state across separate parent invocations. This pattern commonly
    appears when an agent is invoked from inside a tool used by another agent.
    """

    # Build a subgraph that interrupts before echoing.
    # Two nodes: "process" interrupts then echoes, "respond" returns "Done".
    def process(state: MessagesState) -> dict:
        interrupt("continue?")
        return {
            "messages": [AIMessage(content=f"Processing: {state['messages'][-1].text}")]
        }

    def respond(state: MessagesState) -> dict:
        return {"messages": [AIMessage(content="Done")]}

    inner = (
        StateGraph(MessagesState)
        .add_node("process", process)
        .add_node("respond", respond)
        .add_edge(START, "process")
        .add_edge("process", "respond")
        .compile()
    )

    def call_inner(state: ParentState) -> dict:
        resp = inner.invoke({"messages": [HumanMessage(content="apples")]})
        return {"result": resp["messages"][-1].text}

    parent = (
        StateGraph(ParentState)
        .add_node("call_inner", call_inner)
        .add_edge(START, "call_inner")
        .compile(checkpointer=sync_checkpointer)
    )
    config = {"configurable": {"thread_id": str(uuid4())}}

    # First invoke hits the interrupt
    result = parent.invoke({"result": ""}, config)
    assert result == {
        "result": "",
        "__interrupt__": [Interrupt(value="continue?", id=AnyStr())],
    }

    # Resume completes the subgraph
    result = parent.invoke(Command(resume=True), config)
    assert result == {"result": "Done"}


def test_stateless_state_resets(
    sync_checkpointer: BaseCheckpointSaver,
) -> None:
    """Tests that a subgraph compiled with checkpointer=None (the default) does
    not retain any message history between separate parent invocations. Each time
    the parent graph invokes the subgraph, it starts with a clean slate. This
    confirms the "stateless" behavior: even though the parent has a checkpointer,
    the subgraph state is not persisted across calls.
    """

    # Build a simple echo subgraph: echoes "Processing: <input>"
    def echo(state: MessagesState) -> dict:
        return {
            "messages": [AIMessage(content=f"Processing: {state['messages'][-1].text}")]
        }

    inner = (
        StateGraph(MessagesState)
        .add_node("echo", echo)
        .add_edge(START, "echo")
        .compile()
    )

    subgraph_messages: list[list[str]] = []
    call_count = 0

    def call_inner(state: ParentState) -> dict:
        nonlocal call_count
        call_count += 1
        topic = "apples" if call_count == 1 else "bananas"
        resp = inner.invoke(
            {"messages": [HumanMessage(content=f"tell me about {topic}")]}
        )
        subgraph_messages.append([m.text for m in resp["messages"]])
        return {"result": resp["messages"][-1].text}

    parent = (
        StateGraph(ParentState)
        .add_node("call_inner", call_inner)
        .add_edge(START, "call_inner")
        .compile(checkpointer=sync_checkpointer)
    )
    config = {"configurable": {"thread_id": str(uuid4())}}

    result1 = parent.invoke({"result": ""}, config)
    assert result1 == {"result": "Processing: tell me about apples"}

    result2 = parent.invoke({"result": ""}, config)
    assert result2 == {"result": "Processing: tell me about bananas"}

    # Both invocations produce fresh history — no memory of prior call
    assert subgraph_messages[0] == [
        "tell me about apples",
        "Processing: tell me about apples",
    ]
    assert subgraph_messages[1] == [
        "tell me about bananas",
        "Processing: tell me about bananas",
    ]


def test_stateless_state_resets_with_interrupt(
    sync_checkpointer: BaseCheckpointSaver,
) -> None:
    """Tests that a subgraph compiled with checkpointer=None resets its state
    between parent invocations even when interrupt/resume is used. The subgraph
    is invoked twice from the parent, each time with an interrupt that must be
    resumed. After both invoke+resume cycles, each subgraph run should only
    contain its own messages — no bleed-over from the previous run.
    """

    # Build a subgraph that interrupts before echoing, then responds "Done"
    def process(state: MessagesState) -> dict:
        interrupt("continue?")
        return {
            "messages": [AIMessage(content=f"Processing: {state['messages'][-1].text}")]
        }

    def respond(state: MessagesState) -> dict:
        return {"messages": [AIMessage(content="Done")]}

    inner = (
        StateGraph(MessagesState)
        .add_node("process", process)
        .add_node("respond", respond)
        .add_edge(START, "process")
        .add_edge("process", "respond")
        .compile()
    )

    subgraph_messages: list[list[str]] = []
    call_count = 0

    def call_inner(state: ParentState) -> dict:
        nonlocal call_count
        call_count += 1
        topic = "apples" if call_count == 1 else "bananas"
        resp = inner.invoke(
            {"messages": [HumanMessage(content=f"tell me about {topic}")]}
        )
        subgraph_messages.append([m.text for m in resp["messages"]])
        return {"result": resp["messages"][-1].text}

    parent = (
        StateGraph(ParentState)
        .add_node("call_inner", call_inner)
        .add_edge(START, "call_inner")
        .compile(checkpointer=sync_checkpointer)
    )
    config = {"configurable": {"thread_id": str(uuid4())}}

    # First invoke+resume cycle
    result = parent.invoke({"result": ""}, config)
    assert result == {
        "result": "",
        "__interrupt__": [Interrupt(value="continue?", id=AnyStr())],
    }
    result = parent.invoke(Command(resume=True), config)
    assert result == {"result": "Done"}

    # Second invoke+resume cycle
    result = parent.invoke({"result": ""}, config)
    assert result == {
        "result": "",
        "__interrupt__": [Interrupt(value="continue?", id=AnyStr())],
    }
    result = parent.invoke(Command(resume=True), config)
    assert result == {"result": "Done"}

    # Both invocations produce fresh history — no memory of prior call
    assert subgraph_messages[0] == [
        "tell me about apples",
        "Processing: tell me about apples",
        "Done",
    ]
    assert subgraph_messages[1] == [
        "tell me about bananas",
        "Processing: tell me about bananas",
        "Done",
    ]


# -- checkpointer=False --


def test_checkpointer_false_no_persistence(
    sync_checkpointer: BaseCheckpointSaver,
) -> None:
    """Tests that a subgraph compiled with checkpointer=False gets no
    persistence at all, even when the parent graph has a checkpointer. Unlike
    the default (checkpointer=None) which inherits just enough from the parent
    to support interrupt/resume, checkpointer=False explicitly opts out of all
    checkpoint behavior. Each invocation starts completely fresh.
    """

    # Build a simple echo subgraph with checkpointer=False
    def echo(state: MessagesState) -> dict:
        return {
            "messages": [AIMessage(content=f"Processed: {state['messages'][-1].text}")]
        }

    inner = (
        StateGraph(MessagesState)
        .add_node("echo", echo)
        .add_edge(START, "echo")
        .compile(checkpointer=False)
    )

    subgraph_messages: list[list[str]] = []
    call_count = 0

    def call_inner(state: ParentState) -> dict:
        nonlocal call_count
        call_count += 1
        topic = "apples" if call_count == 1 else "bananas"
        resp = inner.invoke(
            {"messages": [HumanMessage(content=f"tell me about {topic}")]}
        )
        subgraph_messages.append([m.text for m in resp["messages"]])
        return {"result": resp["messages"][-1].text}

    parent = (
        StateGraph(ParentState)
        .add_node("call_inner", call_inner)
        .add_edge(START, "call_inner")
        .compile(checkpointer=sync_checkpointer)
    )
    config = {"configurable": {"thread_id": str(uuid4())}}

    result1 = parent.invoke({"result": ""}, config)
    assert result1 == {"result": "Processed: tell me about apples"}

    result2 = parent.invoke({"result": ""}, config)
    assert result2 == {"result": "Processed: tell me about bananas"}

    # Both start fresh — no history from first call
    assert subgraph_messages[0] == [
        "tell me about apples",
        "Processed: tell me about apples",
    ]
    assert subgraph_messages[1] == [
        "tell me about bananas",
        "Processed: tell me about bananas",
    ]


# -- checkpointer=True (stateful) --


def test_stateful_state_accumulates(
    sync_checkpointer: BaseCheckpointSaver,
) -> None:
    """Tests that a subgraph compiled with checkpointer=True ("stateful")
    retains its message history across separate parent invocations. To enable
    this, the subgraph is wrapped in an outer graph compiled with
    checkpointer=True — this wrapper gives the inner subgraph its own persistent
    checkpoint namespace. After two parent calls, the second subgraph invocation
    should see messages from both the first and second calls.
    """

    # Build a simple echo subgraph
    def echo(state: MessagesState) -> dict:
        return {
            "messages": [AIMessage(content=f"Processing: {state['messages'][-1].text}")]
        }

    inner = (
        StateGraph(MessagesState)
        .add_node("echo", echo)
        .add_edge(START, "echo")
        .compile()
    )

    # Wrap the inner subgraph with checkpointer=True to enable stateful.
    # The wrapper graph gives the subgraph its own persistent checkpoint
    # namespace, keyed by the node name ("agent").
    wrapper = (
        StateGraph(MessagesState)
        .add_node("agent", inner)
        .add_edge(START, "agent")
        .compile(checkpointer=True)
    )

    subgraph_messages: list[list[str]] = []
    topics = ["apples", "bananas"]

    def call_inner(state: ParentState) -> dict:
        topic = topics[len(subgraph_messages)]
        resp = wrapper.invoke(
            {"messages": [HumanMessage(content=f"tell me about {topic}")]}
        )
        subgraph_messages.append([m.text for m in resp["messages"]])
        return {"result": resp["messages"][-1].text}

    parent = (
        StateGraph(ParentState)
        .add_node("call_inner", call_inner)
        .add_edge(START, "call_inner")
        .compile(checkpointer=sync_checkpointer)
    )
    config = {"configurable": {"thread_id": str(uuid4())}}

    result1 = parent.invoke({"result": ""}, config)
    assert result1 == {"result": "Processing: tell me about apples"}

    result2 = parent.invoke({"result": ""}, config)
    assert result2 == {"result": "Processing: tell me about bananas"}

    # First call: fresh history
    assert subgraph_messages[0] == [
        "tell me about apples",
        "Processing: tell me about apples",
    ]
    # Second call: retains messages from first call
    assert subgraph_messages[1] == [
        "tell me about apples",
        "Processing: tell me about apples",
        "tell me about bananas",
        "Processing: tell me about bananas",
    ]


def test_stateful_state_accumulates_with_interrupt(
    sync_checkpointer: BaseCheckpointSaver,
) -> None:
    """Tests that a stateful subgraph (checkpointer=True) retains its
    message history across parent invocations even when interrupt/resume is
    involved. The subgraph interrupts before echoing, then responds "Done".
    After two invoke+resume cycles, the second run should contain the full
    accumulated history from both calls.
    """

    # Build a subgraph that interrupts before echoing, then responds "Done"
    def process(state: MessagesState) -> dict:
        interrupt("continue?")
        return {
            "messages": [AIMessage(content=f"Processing: {state['messages'][-1].text}")]
        }

    def respond(state: MessagesState) -> dict:
        return {"messages": [AIMessage(content="Done")]}

    inner = (
        StateGraph(MessagesState)
        .add_node("process", process)
        .add_node("respond", respond)
        .add_edge(START, "process")
        .add_edge("process", "respond")
        .compile()
    )

    # Wrap with checkpointer=True for stateful
    wrapper = (
        StateGraph(MessagesState)
        .add_node("agent", inner)
        .add_edge(START, "agent")
        .compile(checkpointer=True)
    )

    subgraph_messages: list[list[str]] = []
    topics = ["apples", "bananas"]

    def call_inner(state: ParentState) -> dict:
        topic = topics[len(subgraph_messages)]
        resp = wrapper.invoke(
            {"messages": [HumanMessage(content=f"tell me about {topic}")]}
        )
        subgraph_messages.append([m.text for m in resp["messages"]])
        return {"result": resp["messages"][-1].text}

    parent = (
        StateGraph(ParentState)
        .add_node("call_inner", call_inner)
        .add_edge(START, "call_inner")
        .compile(checkpointer=sync_checkpointer)
    )
    config = {"configurable": {"thread_id": str(uuid4())}}

    # First invoke+resume cycle
    result = parent.invoke({"result": ""}, config)
    assert result == {
        "result": "",
        "__interrupt__": [Interrupt(value="continue?", id=AnyStr())],
    }
    result = parent.invoke(Command(resume=True), config)
    assert result == {"result": "Done"}

    # Second invoke+resume cycle
    result = parent.invoke({"result": ""}, config)
    assert result == {
        "result": "",
        "__interrupt__": [Interrupt(value="continue?", id=AnyStr())],
    }
    result = parent.invoke(Command(resume=True), config)
    assert result == {"result": "Done"}

    # First call: fresh history
    assert subgraph_messages[0] == [
        "tell me about apples",
        "Processing: tell me about apples",
        "Done",
    ]
    # Second call: retains messages from first call
    assert subgraph_messages[1] == [
        "tell me about apples",
        "Processing: tell me about apples",
        "Done",
        "tell me about bananas",
        "Processing: tell me about bananas",
        "Done",
    ]


def test_stateful_interrupt_resume(
    sync_checkpointer: BaseCheckpointSaver,
) -> None:
    """Tests that a stateful subgraph (checkpointer=True) correctly
    supports interrupt/resume while also accumulating state. Each invoke+resume
    pair triggers the subgraph, and after the second pair completes we verify
    both the per-step invoke outputs and the accumulated message history. This
    exercises the full lifecycle: interrupt, resume, state accumulation.
    """

    # Build a subgraph that interrupts before echoing, then responds "Done"
    def process(state: MessagesState) -> dict:
        interrupt("continue?")
        return {
            "messages": [AIMessage(content=f"Processing: {state['messages'][-1].text}")]
        }

    def respond(state: MessagesState) -> dict:
        return {"messages": [AIMessage(content="Done")]}

    inner = (
        StateGraph(MessagesState)
        .add_node("process", process)
        .add_node("respond", respond)
        .add_edge(START, "process")
        .add_edge("process", "respond")
        .compile()
    )

    # Wrap with checkpointer=True for stateful
    wrapper = (
        StateGraph(MessagesState)
        .add_node("agent", inner)
        .add_edge(START, "agent")
        .compile(checkpointer=True)
    )

    subgraph_messages: list[list[str]] = []
    topics = ["apples", "bananas"]

    def call_inner(state: ParentState) -> dict:
        topic = topics[len(subgraph_messages)]
        resp = wrapper.invoke(
            {"messages": [HumanMessage(content=f"tell me about {topic}")]}
        )
        subgraph_messages.append([m.text for m in resp["messages"]])
        return {"result": resp["messages"][-1].text}

    parent = (
        StateGraph(ParentState)
        .add_node("call_inner", call_inner)
        .add_edge(START, "call_inner")
        .compile(checkpointer=sync_checkpointer)
    )
    config = {"configurable": {"thread_id": str(uuid4())}}

    # First invocation: hits interrupt
    result = parent.invoke({"result": ""}, config)
    assert result == {
        "result": "",
        "__interrupt__": [Interrupt(value="continue?", id=AnyStr())],
    }

    # Resume: completes first call
    result = parent.invoke(Command(resume=True), config)
    assert result == {"result": "Done"}
    assert subgraph_messages[0] == [
        "tell me about apples",
        "Processing: tell me about apples",
        "Done",
    ]

    # Second invocation: hits interrupt, state accumulated from first call
    result = parent.invoke({"result": ""}, config)
    assert result == {
        "result": "",
        "__interrupt__": [Interrupt(value="continue?", id=AnyStr())],
    }

    # Resume: completes second call with accumulated state
    result = parent.invoke(Command(resume=True), config)
    assert result == {"result": "Done"}
    assert subgraph_messages[1] == [
        "tell me about apples",
        "Processing: tell me about apples",
        "Done",
        "tell me about bananas",
        "Processing: tell me about bananas",
        "Done",
    ]


def test_stateful_namespace_isolation(
    sync_checkpointer: BaseCheckpointSaver,
) -> None:
    """Tests that two different stateful subgraphs (checkpointer=True)
    maintain completely independent state when they use different wrapper node
    names. A "fruit_agent" and "veggie_agent" are each wrapped in their own
    stateful graph. After two parent invocations, each agent should only
    see its own accumulated history with no cross-contamination between them.
    """

    # Build two simple echo subgraphs with different prefixes
    def fruit_echo(state: MessagesState) -> dict:
        return {"messages": [AIMessage(content=f"Fruit: {state['messages'][-1].text}")]}

    def veggie_echo(state: MessagesState) -> dict:
        return {
            "messages": [AIMessage(content=f"Veggie: {state['messages'][-1].text}")]
        }

    fruit_inner = (
        StateGraph(MessagesState)
        .add_node("echo", fruit_echo)
        .add_edge(START, "echo")
        .compile()
    )
    veggie_inner = (
        StateGraph(MessagesState)
        .add_node("echo", veggie_echo)
        .add_edge(START, "echo")
        .compile()
    )

    # Wrap each with checkpointer=True, using different node names to get
    # independent checkpoint namespaces
    fruit = (
        StateGraph(MessagesState)
        .add_node("fruit_agent", fruit_inner)
        .add_edge(START, "fruit_agent")
        .compile(checkpointer=True)
    )
    veggie = (
        StateGraph(MessagesState)
        .add_node("veggie_agent", veggie_inner)
        .add_edge(START, "veggie_agent")
        .compile(checkpointer=True)
    )

    fruit_msgs: list[list[str]] = []
    veggie_msgs: list[list[str]] = []
    call_count = 0

    def call_both(state: ParentState) -> dict:
        nonlocal call_count
        call_count += 1
        suffix = "round 1" if call_count == 1 else "round 2"
        f = fruit.invoke({"messages": [HumanMessage(content=f"cherries {suffix}")]})
        v = veggie.invoke({"messages": [HumanMessage(content=f"broccoli {suffix}")]})
        fruit_msgs.append([m.text for m in f["messages"]])
        veggie_msgs.append([m.text for m in v["messages"]])
        return {"result": f["messages"][-1].text}

    parent = (
        StateGraph(ParentState)
        .add_node("call_both", call_both)
        .add_edge(START, "call_both")
        .compile(checkpointer=sync_checkpointer)
    )
    config = {"configurable": {"thread_id": str(uuid4())}}

    result1 = parent.invoke({"result": ""}, config)
    assert result1 == {"result": "Fruit: cherries round 1"}

    result2 = parent.invoke({"result": ""}, config)
    assert result2 == {"result": "Fruit: cherries round 2"}

    # First call: each agent sees only its own history
    assert fruit_msgs[0] == ["cherries round 1", "Fruit: cherries round 1"]
    assert veggie_msgs[0] == ["broccoli round 1", "Veggie: broccoli round 1"]

    # Second call: each accumulated independently — no cross-contamination
    assert fruit_msgs[1] == [
        "cherries round 1",
        "Fruit: cherries round 1",
        "cherries round 2",
        "Fruit: cherries round 2",
    ]
    assert veggie_msgs[1] == [
        "broccoli round 1",
        "Veggie: broccoli round 1",
        "broccoli round 2",
        "Veggie: broccoli round 2",
    ]


def test_child_with_own_thread_id_keeps_namespace(
    sync_checkpointer: BaseCheckpointSaver,
) -> None:
    """A child graph invoked from inside a parent node with its own thread_id
    must store and read its checkpoint under its own namespace, not inherit the
    parent task's checkpoint_ns.
    """

    class ChildState(TypedDict):
        count: int

    def child_node(state: ChildState) -> dict:
        return {"count": (state.get("count") or 0) + 1}

    child = (
        StateGraph(ChildState)
        .add_node("n", child_node)
        .add_edge(START, "n")
        .compile(checkpointer=sync_checkpointer)
    )

    child_thread = str(uuid4())
    child_config = {"configurable": {"thread_id": child_thread}}

    def parent_node(state: ParentState) -> dict:
        child.invoke({}, config=child_config)
        return {"result": "ok"}

    parent = (
        StateGraph(ParentState)
        .add_node("p", parent_node)
        .add_edge(START, "p")
        .compile(checkpointer=sync_checkpointer)
    )
    parent_config = {"configurable": {"thread_id": str(uuid4())}}

    parent.invoke({"result": ""}, config=parent_config)
    state1 = child.get_state(child_config)
    assert state1.values.get("count") == 1
    assert state1.config["configurable"]["checkpoint_ns"] == ""

    parent.invoke({"result": ""}, config=parent_config)
    state2 = child.get_state(child_config)
    assert state2.values.get("count") == 2