from typing import Literal
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import tools_condition
from agent.nodes import call_model, tool_node
from agent.state import State
from langchain_core.messages import AIMessage, HumanMessage


def get_agent():
    """
    Get our LangGraph agent with the given model and tools.
    """
   
    class GraphConfig(TypedDict):
        app_name: str

    graph = StateGraph(State, context_schema=GraphConfig)
    
    # Add nodes
    graph.add_node("agent", call_model)
    graph.add_node("tools", tool_node)
    
    graph.add_edge(START, "agent")
    graph.add_conditional_edges("agent", tools_condition)
    graph.add_edge("tools", "agent")
        
    return graph.compile()

# test
if __name__ == "__main__":
    import os
    
    question = "When was Andrej Karpathy born?"
    messages = [HumanMessage(content=question)]
    
    try:
        graph = get_agent()
        from agent.config import create_agent_config
        config = create_agent_config("OracleBot")
        result = graph.invoke({"messages": messages}, config)
        print("\n=== Agent Response ===")
        for m in result["messages"]:
            m.pretty_print()
    except Exception as e:
        print(f"Error running agent: {e}")
        print("Make sure your API key is correct and the service is accessible.")