#!/usr/bin/env python3
"""
agent_graph_demo.py: a working orchestrator-worker agent graph in LangGraph,
and the one edge that quietly breaks it.

Companion to the walkthrough "The Graph Looks Right. The Merge Is Where It Breaks."

    pip install langgraph langchain-core
    python agent_graph_demo.py            # the SAFE graph, with a human gate
    python agent_graph_demo.py --trap     # the SAME graph wired the unsafe way

There is no API key and no network here. `call_model` is a deterministic stub
so the graph's SHAPE is what you are watching, not a model's mood. The last
function in this file shows the two lines you change to use a real model.

Task: three read-only workers each check one source about a release (the test
run, the error metrics, the changelog). A single writer reads all three and
makes ONE call: ship, hold, or roll back.
"""
import argparse
import operator
from typing import Annotated, List, Literal, TypedDict

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import interrupt, Command

CALLS = {"n": 0}   # a crude cost meter: how many model calls the graph made


def call_model(role: str, source: str) -> str:
    """A stand-in for a real LLM. Deterministic so the graph is what varies."""
    CALLS["n"] += 1
    verdict = {
        "tests": "hold",        # a flaky test failed
        "metrics": "ship it",   # error rate is flat
        "changelog": "roll back",  # an undocumented breaking change
    }[source]
    return verdict


# ---------------------------------------------------------------------------
# THE SAFE GRAPH
# GATHER state has many writers (workers append structured findings). DECISION
# state has one owner (the writer applies an explicit policy). Different jobs,
# different channels.
# ---------------------------------------------------------------------------
Verdict = Literal["ship it", "hold", "roll back"]


class Finding(TypedDict):
    reviewer: str
    verdict: Verdict


class State(TypedDict):
    findings: Annotated[List[Finding], operator.add]   # gather: many writers append
    decision: Verdict                                  # decision: one owner writes


def orchestrator(state: State) -> dict:
    # In a real build this would PLAN the workers from the question. Here the
    # plan is fixed so you can see the topology. It writes nothing a worker writes.
    return {}


def make_worker(source: str):
    def worker(state: State) -> dict:
        verdict = call_model(role="reviewer", source=source)
        # Each worker emits its OWN structured finding. It never writes `decision`.
        return {"findings": [{"reviewer": source, "verdict": verdict}]}
    return worker


def decide(verdicts: List[Verdict]) -> Verdict:
    # The decision POLICY, named and owned. This one is any-veto by severity;
    # swap it for majority / unanimous / weighted / threshold as your gate needs.
    if "roll back" in verdicts:
        return "roll back"
    if "hold" in verdicts:
        return "hold"
    return "ship it"


EXPECTED = {"tests", "metrics", "changelog"}


def writer(state: State) -> dict:
    # The single owner reads ALL findings and applies the policy. It is the only
    # node that writes `decision`. Ownership (one writer) is only part of the gate:
    # completeness (every reviewer reported) and uniqueness (each reported exactly
    # once) matter too, or a missing verdict silently ships and a duplicate one can
    # tip a majority rule. Fail closed on either.
    reviewers = [f["reviewer"] for f in state["findings"]]
    if sorted(reviewers) != sorted(EXPECTED):   # missing or duplicate -> not a ship
        return {"decision": "hold"}
    return {"decision": decide([f["verdict"] for f in state["findings"]])}


def human_gate(state: State):
    # interrupt() PAUSES the graph and hands control back to you. Resuming
    # re-runs this node FROM THE TOP, so keep side effects below the interrupt.
    approved = interrupt({"proposed": state["decision"],
                          "ask": "approve this decision? (resume with True/False)"})
    if not approved:
        return {"decision": state["decision"] + "  [VETOED by human]"}
    return {"decision": state["decision"] + "  [approved]"}


def build_safe_graph():
    g = StateGraph(State)
    g.add_node("orchestrator", orchestrator)
    g.add_node("writer", writer)
    g.add_node("human_gate", human_gate)
    for src in ("tests", "metrics", "changelog"):
        g.add_node(src, make_worker(src))
        g.add_edge("orchestrator", src)   # fan out (reads run in parallel)
        g.add_edge(src, "writer")         # fan in to the one writer
    g.add_edge(START, "orchestrator")
    g.add_edge("writer", "human_gate")
    g.add_edge("human_gate", END)
    # A checkpointer is REQUIRED for interrupt(): no saved state, nothing to resume.
    return g.compile(checkpointer=MemorySaver())


def run_safe():
    app = build_safe_graph()
    cfg = {"configurable": {"thread_id": "release-1"}}
    result = app.invoke({"question": "Do we ship?", "findings": [], "decision": ""}, cfg)
    # The run paused at the human gate. `result` carries the interrupt payload.
    payload = result["__interrupt__"][0].value
    print("PAUSED at human gate. Proposed:", payload["proposed"])
    final = app.invoke(Command(resume=False), cfg)   # human vetoes
    print("FINAL decision:", final["decision"])
    print("model calls made:", CALLS["n"])
    print("\nSAFE: three workers, three findings, ONE writer. 3 of 3 findings"
          " reached the decision, and a human held the veto.")


# ---------------------------------------------------------------------------
# THE TRAP
# The SAME three workers, but each writes `decision` directly. `decision` is a
# DECISION channel (one answer should win), not a gather channel. Watch the
# framework refuse, then watch the reducer you add to silence it hide the conflict.
# ---------------------------------------------------------------------------
class TrapState(TypedDict):
    decision: str   # a plain DECISION channel; three parallel writers collide here


class TrapStateReducer(TypedDict):
    # The reducer a hurried dev adds to make the error stop: keep the latest write.
    # It silences the error and silently drops the losers.
    decision: Annotated[str, lambda old, new: new]


def run_trap():
    # Step 1: the naive version. The framework STOPS you, loudly.
    g = StateGraph(TrapState)
    g.add_node("orchestrator", lambda s: {})
    for src in ("tests", "metrics", "changelog"):
        g.add_node(src, (lambda src=src: (lambda s: {"decision": call_model("r", src)}))())
        g.add_edge("orchestrator", src); g.add_edge(src, END)
    g.add_edge(START, "orchestrator")
    try:
        g.compile().invoke({"decision": ""})
    except Exception as e:
        print("STEP 1 (plain channel): the framework refused ->",
              type(e).__name__ + ":", str(e).splitlines()[0])

    # Step 2: you take the error's advice and add a reducer. Now it "works".
    g2 = StateGraph(TrapStateReducer)
    g2.add_node("orchestrator", lambda s: {})
    for src in ("tests", "metrics", "changelog"):
        g2.add_node(src, (lambda src=src: (lambda s: {"decision": call_model("r", src)}))())
        g2.add_edge("orchestrator", src); g2.add_edge(src, END)
    g2.add_edge(START, "orchestrator")
    kept = g2.compile().invoke({"decision": ""})["decision"]
    print("STEP 2 (reducer added): no error. the decision channel now holds:", repr(kept))
    print("        the other two reviewers' verdicts were silently dropped")
    print("\nTRAP: the error could not tell a decision channel from a gather channel, and"
          " neither could the reducer you added. 1 of 3 verdicts survives; two vanish, no error.")
    print("      (an operator.add reducer here would instead change the channel from a string"
          " to a list, and a consumer expecting a string mishandles it just as quietly.)")


def use_a_real_model_instead():
    """Not called. The only change from the stub to a real model:

        from langchain_anthropic import ChatAnthropic     # pip install langchain-anthropic
        _llm = ChatAnthropic(model="claude-sonnet-4-5")    # reads ANTHROPIC_API_KEY

    then inside a worker, replace the stub line with:

        verdict = _llm.invoke(f"Review the {source} and answer ship it/hold/roll back").content

    The graph, the channels, and the read/write rule do not change. The model
    is a node; the safety is in the wiring.
    """


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--trap", action="store_true", help="wire the same graph the unsafe way")
    args = ap.parse_args()
    run_trap() if args.trap else run_safe()