"""Actual LangGraph execution around the deterministic Engine; no LLM.

Optional dependency: langgraph==1.2.16. Both checkpoint stores are in memory.
"""
from typing import TypedDict

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

from engine import BaselinePlanner, Checkpoints, Engine
from workflow import Request


class GraphState(TypedDict):
    run_id: str
    phase: str


def build_graph(engine, planner=None):
    planner = planner if planner is not None else BaselinePlanner()

    def plan(state):
        result = engine.plan(state["run_id"], planner, now_tick=5)
        return {"phase": result["phase"]}

    def invoke(state):
        return {"phase": engine.invoke(state["run_id"])["phase"]}

    def review(state):
        saved = engine.store.get(state["run_id"])
        # No side effect before interrupt: this function is restarted on resume.
        decision = interrupt({"simulation_only": True, "run_id": saved["id"],
                              "revision": saved["revision"], "draft_digest": saved["draft_digest"],
                              "draft": saved["draft"], "deadline": saved["deadline"]})
        if not isinstance(decision, dict) or set(decision) != {
            "revision", "draft_digest", "reviewer", "approve", "now_tick"
        }:
            raise ValueError("invalid synthetic review payload")
        result = engine.review(saved["id"], decision["revision"], decision["draft_digest"],
                               decision["reviewer"], decision["approve"], decision["now_tick"])
        return {"phase": result["phase"]}

    def finalize(state):
        return {"phase": engine.finalize_simulation(state["run_id"], now_tick=6)["phase"]}

    builder = StateGraph(GraphState)
    for name, node in (("plan", plan), ("invoke", invoke), ("review", review), ("finalize", finalize)):
        builder.add_node(name, node)
    builder.add_edge(START, "plan")
    builder.add_conditional_edges("plan", lambda s: s["phase"],
                                  {"RESERVED": "invoke", "REVIEW": "review", "STOPPED": END})
    builder.add_conditional_edges("invoke", lambda s: s["phase"],
                                  {"ACTIVE": "plan", "STOPPED": END})
    builder.add_conditional_edges("review", lambda s: s["phase"],
                                  {"APPROVED_FOR_SIMULATION": "finalize", "REJECTED": END, "STOPPED": END})
    builder.add_edge("finalize", END)
    return builder.compile(checkpointer=InMemorySaver())


def decision_for(saved, approve=True):
    return {"revision": saved["revision"], "draft_digest": saved["draft_digest"],
            "reviewer": "synthetic-reviewer-a", "approve": approve, "now_tick": 6}


def main():
    store = Checkpoints()
    try:
        engine = Engine(store)
        engine.start(Request("GRAPH-1", "EX-001", 1927))
        graph = build_graph(engine)
        config = {"configurable": {"thread_id": "synthetic-graph-1"}, "recursion_limit": 64}
        paused = graph.invoke({"run_id": "GRAPH-1", "phase": "ACTIVE"}, config)
        print("ACTUAL LANGGRAPH; SCRIPTED CONTROL; NO LLM / REMOTE SERVICE")
        print("Interrupted:", "__interrupt__" in paused, "phase:", store.get("GRAPH-1")["phase"])
        final = graph.invoke(Command(resume=decision_for(store.get("GRAPH-1"))), config)
        print("Resumed phase:", final["phase"], "receipt:", store.get("GRAPH-1")["receipt"])
    finally:
        store.close()


if __name__ == "__main__":
    main()