"""
Basic LangGraph Example - A simple conversational agent with state management.

This demonstrates core LangGraph concepts:
- State schemas
- Nodes (functions that process state)
- Edges (connections between nodes)
- Conditional routing
"""

from typing import TypedDict, Literal, Annotated
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages


# Define the state schema
class State(TypedDict):
    """State represents what flows through our graph."""
    messages: Annotated[list, add_messages]
    counter: int


# Define node functions
def greeting_node(state: State) -> dict:
    """Process greeting messages."""
    print("Processing greeting...")
    return {"counter": state["counter"] + 1}


def math_node(state: State) -> dict:
    """Process math-related messages."""
    print("Processing math request...")
    return {"counter": state["counter"] + 1}


def default_node(state: State) -> dict:
    """Handle unknown/unrecognized messages."""
    print("Handling general message...")
    return {"counter": state["counter"] + 1}


# Define the router function
def route_message(state: State) -> Literal["greeting", "math", "default"]:
    """Route to appropriate node based on message content."""
    messages = state["messages"]
    if not messages:
        return "default"

    last_message = messages[-1].lower() if isinstance(messages[-1], str) else str(messages[-1]).lower()

    if any(word in last_message for word in ["hello", "hi", "hey", "greetings"]):
        return "greeting"
    elif any(word in last_message for word in ["add", "subtract", "multiply", "divide", "calculate"]):
        return "math"
    else:
        return "default"


# Build the graph
def create_graph():
    """Create and compile the LangGraph workflow."""

    # Initialize the graph builder
    builder = StateGraph(State)

    # Add nodes
    builder.add_node("greeting", greeting_node)
    builder.add_node("math", math_node)
    builder.add_node("default", default_node)

    # Add edges with conditional routing
    builder.add_conditional_edges(
        START,
        route_message,
        {
            "greeting": "greeting",
            "math": "math",
            "default": "default"
        }
    )

    # All nodes lead to END (or could chain to other nodes)
    builder.add_edge("greeting", END)
    builder.add_edge("math", END)
    builder.add_edge("default", END)

    # Compile the graph
    return builder.compile()


def main():
    """Run the basic LangGraph agent."""
    graph = create_graph()

    print("=== Basic LangGraph Agent ===")
    print("Try messages like: 'hello', 'add 5 and 3', or anything else")
    print()

    # Test cases
    test_messages = [
        "Hello there!",
        "Can you add 5 and 3?",
        "What's the weather like?"
    ]

    for message in test_messages:
        print(f"Input: {message}")
        initial_state = {"messages": [message], "counter": 0}
        result = graph.invoke(initial_state)
        print(f"Output: Processed with counter = {result['counter']}")
        print()


if __name__ == "__main__":
    main()