"""
Demonstrate compiling a LangGraph and visualizing its structure.
Requires no external viz packages for ASCII. To export PNG, install:
  uv pip install -e '.[viz]'
and ensure Graphviz is present on your system.
"""

from typing import TypedDict, Annotated
import operator

from langgraph.graph import StateGraph, START, END


class S(TypedDict):
    a: int
    b: int
    c: int


def n1(s: S) -> S:
    return {"a": s.get("a", 0) + 1, "b": s.get("b", 0), "c": s.get("c", 0)}


def n2(s: S) -> S:
    return {"a": s.get("a", 0), "b": s.get("b", 0) + 2, "c": s.get("c", 0)}


def n3(s: S) -> S:
    return {"a": s.get("a", 0), "b": s.get("b", 0), "c": s.get("c", 0) + 3}


def build_app():
    g = StateGraph(S)
    g.add_node("n1", n1)
    g.add_node("n2", n2)
    g.add_node("n3", n3)
    g.add_edge(START, "n1")
    g.add_edge("n1", "n2")
    g.add_edge("n2", "n3")
    g.add_edge("n3", END)
    return g.compile()


def main():
    app = build_app()
    graph = app.get_graph()

    # Try ASCII render if available
    draw_ascii = getattr(graph, "draw_ascii", None)
    if callable(draw_ascii):
        print("=== Graph ASCII ===")
        print(draw_ascii())

    # Try PNG export if available and pydot/graphviz installed
    draw_png = getattr(graph, "draw_png", None)
    if callable(draw_png):
        try:
            out = "graph_demo.png"
            draw_png(out)
            print(f"Saved PNG: {out}")
        except Exception as e:
            print(f"PNG export unavailable: {e}")


if __name__ == "__main__":
    main()