"""Simple base LangGraph for conversational agent."""
from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_core.language_models.chat_models import BaseChatModel
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages


class AgentState(TypedDict):
    """State for the simple conversational agent."""
    messages: Annotated[Sequence[BaseMessage], add_messages]


def create_simple_graph(llm: BaseChatModel):
    """
    Create a simple conversational graph with LangGraph.
    
    This is a basic graph that takes a message, sends it to the LLM,
    and returns the response. It can be easily replaced with more complex graphs.
    
    Args:
        llm: Language model to use for generation
        
    Returns:
        Compiled LangGraph
    """
    
    def call_model(state: AgentState) -> AgentState:
        """Call the LLM with the current messages."""
        print(f"Calling model with messages: {state['messages']}")
        
        messages = state["messages"]
        response = llm.invoke(messages)
        return {"messages": [response]}
    
    # Build the graph
    workflow = StateGraph(AgentState)
    
    # Add nodes
    workflow.add_node("agent", call_model)
    
    # Set entry point
    workflow.set_entry_point("agent")
    
    # Add edge to end
    workflow.add_edge("agent", END)
    
    # Compile and return
    return workflow.compile()


def create_simple_graph_with_history(llm: BaseChatModel):
    """
    Create a simple conversational graph with conversation history support.
    
    Args:
        llm: Language model to use for generation
        
    Returns:
        Compiled LangGraph
    """
    
    def call_model_with_history(state: AgentState) -> AgentState:
        """Call the LLM with full conversation history."""
        messages = state["messages"]
        response = llm.invoke(messages)
        return {"messages": [response]}
    
    # Build the graph
    workflow = StateGraph(AgentState)
    
    # Add nodes
    workflow.add_node("agent", call_model_with_history)
    
    # Set entry point
    workflow.set_entry_point("agent")
    
    # Add edge to end
    workflow.add_edge("agent", END)
    
    # Compile and return
    return workflow.compile()