"""17 - Checkpoint state history, updates, replay, and forks.

This example is completely local: ``InMemorySaver`` stores checkpoints in RAM.
A checkpoint config (thread_id + checkpoint_id) identifies one exact point in time.

Run: python 17_state_history.py
"""

from typing import TypedDict

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


class State(TypedDict):
    value: int
    note: str


def increment(state: State) -> dict:
    """First deterministic step."""
    return {"value": state["value"] + 1, "note": "incremented"}


def double(state: State) -> dict:
    """Second deterministic step."""
    return {"value": state["value"] * 2, "note": "doubled"}


def build_graph():
    builder = StateGraph(State)
    builder.add_node("increment", increment)
    builder.add_node("double", double)
    builder.add_edge(START, "increment")
    builder.add_edge("increment", "double")
    builder.add_edge("double", END)
    return builder.compile(checkpointer=InMemorySaver())


def short_config(config: dict) -> dict:
    """Only display the portable checkpoint identity, not internal metadata."""
    configurable = config["configurable"]
    return {
        "thread_id": configurable["thread_id"],
        "checkpoint_id": configurable.get("checkpoint_id"),
    }


def main() -> None:
    graph = build_graph()
    thread_config = {"configurable": {"thread_id": "history-demo"}}

    print("=== Initial run ===")
    result = graph.invoke({"value": 3, "note": "input"}, thread_config)
    print("result:", result)
    assert result["value"] == 8

    # get_state with only a thread_id returns that thread's latest snapshot.
    latest = graph.get_state(thread_config)
    print("latest values:", latest.values)
    print("latest next:", latest.next)
    print("latest config:", short_config(latest.config))

    # History is returned newest first. Every StateSnapshot has values, next,
    # config, metadata, created_at, parent_config, and tasks.
    history = list(graph.get_state_history(thread_config))
    print("\n=== State history (newest first) ===")
    for index, snapshot in enumerate(history):
        print(
            index,
            "values=", snapshot.values,
            "next=", snapshot.next,
            "checkpoint=", short_config(snapshot.config),
        )

    # Select the checkpoint immediately after `increment`: `double` is pending.
    before_double = next(snapshot for snapshot in history if snapshot.next == ("double",))
    checkpoint_config = before_double.config
    assert before_double.values["value"] == 4

    print("\n=== Replay from an historical checkpoint ===")
    # Passing the snapshot's config and None resumes its pending work. Earlier
    # nodes are not rerun, so only `double` executes.
    replayed = graph.invoke(None, checkpoint_config)
    print("replayed:", replayed)
    assert replayed["value"] == 8

    print("\n=== Fork by updating an historical checkpoint ===")
    # update_state does not mutate history. It creates a new checkpoint derived
    # from checkpoint_config. The returned config identifies the new branch.
    fork_config = graph.update_state(
        checkpoint_config,
        {"value": 10, "note": "human correction"},
    )
    fork_snapshot = graph.get_state(fork_config)
    print("fork before resume:", fork_snapshot.values, "next=", fork_snapshot.next)
    print("fork config:", short_config(fork_config))

    forked = graph.invoke(None, fork_config)
    print("forked result:", forked)
    assert forked["value"] == 20

    # The old exact checkpoint still contains 4; checkpoint configs make
    # time-travel explicit even after the thread acquires newer branches.
    old_snapshot = graph.get_state(checkpoint_config)
    assert old_snapshot.values["value"] == 4
    print("original historical value is still:", old_snapshot.values["value"])
    print("\nAll state-history demonstrations passed.")


if __name__ == "__main__":
    main()