"""Chapter 09 + 10 — a state-machine engine, the way LangGraph works underneath.

Ch 09 (linear): State (a shared dict), Nodes (state -> partial update), Edges
(who runs next). Ch 10 adds the piece that makes it an *agent* engine rather
than a pipeline: conditional edges (a router picks the next node) and therefore
cycles — so a recursion limit becomes mandatory.
"""

END = "__end__"


class StateGraph:
    def __init__(self):
        self.nodes = {}        # name -> fn(state) -> dict
        self.edges = {}        # name -> next name              (static)
        self.branches = {}     # name -> (router_fn, mapping)   (conditional)
        self.reducers = {}     # key -> fn(old, new)            (Ch 14)
        self.entry = None

    def add_node(self, name, fn):
        self.nodes[name] = fn
        return self

    def add_reducer(self, key, fn):
        """Ch 14: combine (not overwrite) updates to a state key — e.g. append
        messages instead of replacing the list."""
        self.reducers[key] = fn
        return self

    def add_edge(self, src, dst):
        self.edges[src] = dst
        return self

    def add_conditional_edges(self, src, router, mapping):
        """router(state) returns a key; mapping[key] is the next node."""
        self.branches[src] = (router, mapping)
        return self

    def set_entry(self, name):
        self.entry = name
        return self

    def compile(self):
        assert self.entry, "no entry point set"
        return CompiledGraph(self)


class CompiledGraph:
    def __init__(self, g):
        self.g = g

    def _next(self, node, state):
        if node in self.g.branches:
            router, mapping = self.g.branches[node]
            return mapping[router(state)]
        return self.g.edges.get(node, END)

    def invoke(self, state, recursion_limit=25, trace=False,
               checkpointer=None, thread_id=None):
        # Ch 13: carry this thread's remembered state into the new turn.
        if checkpointer and thread_id:
            state = {**checkpointer.load(thread_id), **state}
        else:
            state = dict(state)
        node = self.g.entry
        for _ in range(recursion_limit):
            if node == END:
                break
            update = self.g.nodes[node](state) or {}
            for key, value in update.items():         # Ch 14: reducer-aware merge
                if key in self.g.reducers:
                    state[key] = self.g.reducers[key](state.get(key), value)
                else:
                    state[key] = value
            if trace:
                print(f"  [{node}] -> {update}")
            node = self._next(node, state)
            if checkpointer and thread_id:            # snapshot after each step
                checkpointer.save(thread_id, state)
        else:
            raise RuntimeError(f"recursion limit ({recursion_limit}) hit — cycle?")
        return state


if __name__ == "__main__":
    # A cyclic graph: count up to 3, looping through the same node.
    def tick(state):    return {"n": state["n"] + 1}
    def route(state):   return "loop" if state["n"] < 3 else "done"

    g = StateGraph()
    g.add_node("tick", tick).set_entry("tick")
    g.add_conditional_edges("tick", route, {"loop": "tick", "done": END})
    print("final:", g.compile().invoke({"n": 0}, trace=True))