# file: 02_basic_agent.py
import os
from dotenv import load_dotenv
load_dotenv()

from typing import TypedDict, Annotated, List
from langchain_core.messages import BaseMessage, ToolMessage
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langchain_openai import ChatOpenAI
from tools.word_counter import count_words

# Define the constant for the OpenRouter API URL
OPENROUTER_API_URL = "https://openrouter.ai/api/v1"

# 1. Define the Agent State
class AgentState(TypedDict):
    messages: Annotated[List[BaseMessage], lambda x, y: x + y]

def main():
    """Main function to set up and run a LangGraph agent."""
    if not os.getenv("OPENROUTER_API_KEY"):
        print("Error: OPENROUTER_API_KEY is not set in .env file.")
        return

    print("--- Modern Agent with LangGraph ---")
    
    tools = [count_words]
    
    llm = ChatOpenAI(
        model="anthropic/claude-3.5-haiku",
        temperature=0,
        openai_api_key=os.getenv("OPENROUTER_API_KEY"),
        openai_api_base=OPENROUTER_API_URL,
    )
    
    llm_with_tools = llm.bind_tools(tools)

    # 2. Define the Nodes
    def call_model(state):
        """The node that calls the LLM."""
        messages = state["messages"]
        response = llm_with_tools.invoke(messages)
        return {"messages": [response]}

    tool_node = ToolNode(tools)

    # 3. Define the Conditional Edge
    def should_continue(state):
        """Decides whether to continue or end the loop."""
        last_message = state["messages"][-1]
        if last_message.tool_calls:
            return "continue"
        return "end"

    # 4. Build the Graph
    workflow = StateGraph(AgentState)
    workflow.add_node("agent", call_model)
    workflow.add_node("action", tool_node)
    workflow.set_entry_point("agent")
    workflow.add_conditional_edges(
        "agent",
        should_continue,
        {"continue": "action", "end": END}
    )
    workflow.add_edge("action", "agent")
    app = workflow.compile()

    # 5. Invoke the Agent
    query = "How many words are in the text 'I am a skilled AI engineer'?"
    inputs = {"messages": [("user", query)]}
    
    print(f"Human: {query}")
    
    for chunk in app.stream(inputs, stream_mode="values"):
        last_message = chunk["messages"][-1]
        if isinstance(last_message, ToolMessage):
            print(f"Tool Output: {last_message.content}")

    final_response = app.invoke(inputs)
    final_answer = final_response["messages"][-1].content

    print("\n--- Agent Response ---")
    print(final_answer)
    print("----------------------")

if __name__ == "__main__":
    main()