from langgraph.graph import StateGraph
from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages
from langchain_core.runnables.graph import MermaidDrawMethod

class State(TypedDict):
    messages: Annotated[list[str], add_messages]
    current_step: str

def collect_info(state: State) -> dict:
    print("\n--> In collect_info")
    print(f"Messages before: {state['messages']}")
    
    messages = state["messages"] + ["Information collected"]
    print(f"Messages after: {messages}")
    
    return {
        "messages": messages,
        "current_step": "process"
    }

def process_info(state: State) -> dict:
    print("\n--> In process_info")
    print(f"Messages before: {state['messages']}")
    
    messages = state["messages"] + ["Information processed"]
    print(f"Messages after: {messages}")
    
    return {
        "messages": messages,
        "current_step": "end"
    }

# Create and setup graph
workflow = StateGraph(State)

# Add nodes
workflow.add_node("collect", collect_info)
workflow.add_node("process", process_info)

# Add edges
workflow.add_edge("collect", "process")

# Set entry and finish points
workflow.set_entry_point("collect")
workflow.set_finish_point("process")

app = workflow.compile()

# Run workflow
print("\nStarting workflow...")
initial_state = State(messages=["Starting"], current_step="collect")
final_state = app.invoke(initial_state)
print(f"\nFinal messages: {final_state['messages']}")

# Save the graph visualization as PNG
png_data = app.get_graph().draw_mermaid_png(draw_method=MermaidDrawMethod.API)
with open("workflow_graph.png", "wb") as f:
    f.write(png_data)
print("\nGraph visualization saved as 'workflow_graph.png'")