import asyncio
from datetime import datetime
from pprint import pprint
from dotenv import load_dotenv
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import END, START, StateGraph
from pydantic import BaseModel
from orcheo.graph.state import State
from orcheo.nodes.ai import AgentNode


load_dotenv()


class StructuredResponse(BaseModel):
    """Response model."""

    reply: str
    timestamp: datetime


def build_graph() -> StateGraph:
    """Build the graph."""
    graph = StateGraph(State)
    agent_node = AgentNode(
        name="agent",
        ai_model="openai:gpt-4o-mini",
        predefined_tools=["greet_user"],
        response_format=StructuredResponse,
    )
    graph.add_node("agent", agent_node)
    graph.add_edge(START, "agent")
    graph.add_edge("agent", END)
    return graph


if __name__ == "__main__":
    checkpointer = InMemorySaver()
    workflow = build_graph().compile(checkpointer=checkpointer)
    config = {"configurable": {"thread_id": "123"}}
    result = asyncio.run(
        workflow.ainvoke(
            {
                "messages": [
                    {
                        "role": "user",
                        "content": "Greet user John Doe!",
                    }
                ]
            },
            config,
        )
    )
    pprint(result)
