#!/usr/bin/env python3
"""LockedIn Labs: two Python functions, one real LangGraph workflow.

Python 3.11+, langgraph==1.2.16. Run `test` or `demo` after installation.
Synthetic data only. No model, network, checkpoint, credential or external write.
"""
import json
import sys
from typing import TypedDict

from langgraph.graph import END, START, StateGraph


class DeskState(TypedDict):
    request: str
    source_id: str
    source: str
    draft: str
    status: str


SOURCES = {"P1": "Post the change summary in the release channel."}


def run_case(source_id: str) -> dict:
    observations = []

    def lookup(state: DeskState) -> dict:
        source = SOURCES.get(state["source_id"], "")
        update = {"source": source,
                  "status": "ready" if source else "needs_source"}
        observations.append({"stage": "lookup", "before": dict(state), "update": update})
        return update

    def route(state: DeskState) -> str:
        destination = "draft" if state["source"] else END
        observations.append({"stage": "route", "before": dict(state),
                             "destination": destination})
        return destination

    def draft(state: DeskState) -> dict:
        update = {"draft": f"Suggested reply: {state['source']}", "status": "drafted"}
        observations.append({"stage": "draft", "before": dict(state), "update": update})
        return update

    builder = StateGraph(DeskState)
    builder.add_node("lookup", lookup)
    builder.add_node("draft", draft)
    builder.add_edge(START, "lookup")
    builder.add_conditional_edges("lookup", route, {"draft": "draft", END: END})
    builder.add_edge("draft", END)
    graph = builder.compile()

    initial = {"request": "Where do I post a deployment update?",
               "source_id": source_id, "source": "", "draft": "", "status": "new"}
    result = graph.invoke(initial, {"recursion_limit": 10})

    # Reconstruct only this graph's default, field-replacement updates.
    # Each observed input and the actual final SDK result must agree.
    state = dict(initial)
    frames = [{"stage": "input", "state": dict(state), "update": dict(initial)}]
    for observation in observations:
        assert observation["before"] == state, "SDK input differs from accumulated state"
        state.update(observation.get("update", {}))
        frames.append({"stage": observation["stage"], "state": dict(state),
                       **({"update": observation["update"]} if "update" in observation
                          else {"destination": observation["destination"]})})
    assert result == state, "SDK output differs from accumulated state"
    frames.append({"stage": "end", "state": dict(result)})
    return {"source_id": source_id, "frames": frames}


def test() -> None:
    ordinary = run_case("P1")
    missing = run_case("P404")
    assert [f["stage"] for f in ordinary["frames"]] == ["input", "lookup", "route", "draft", "end"]
    assert [f["stage"] for f in missing["frames"]] == ["input", "lookup", "route", "end"]
    assert ordinary["frames"][2]["destination"] == "draft"
    assert missing["frames"][2]["destination"] == END
    assert ordinary["frames"][-1]["state"]["draft"] == "Suggested reply: " + SOURCES["P1"]
    assert ordinary["frames"][-1]["state"]["status"] == "drafted"
    assert missing["frames"][-1]["state"]["draft"] == ""
    assert missing["frames"][-1]["state"]["status"] == "needs_source"
    for case in [ordinary, missing]:
        assert all(f["state"]["source_id"] == case["source_id"] for f in case["frames"])
        assert case["frames"][1]["state"] == case["frames"][2]["state"], "Routing must not change state"
    assert run_case("P1") == ordinary, "Independent runs should not retain another case's state"
    print("PASS: useful draft, missing-source branch, preserved fields, read-only routing, independent runs")


if __name__ == "__main__":
    if len(sys.argv) != 2 or sys.argv[1] not in {"test", "demo"}:
        raise SystemExit("Usage: python fde-langgraph-first-workflow.py {test|demo}")
    if sys.argv[1] == "test":
        test()
    else:
        print(json.dumps([run_case("P1"), run_case("P404")], indent=2))