from typing import List, TypedDict, Annotated
from langchain_openai import ChatOpenAI
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
from langgraph.graph.message import add_messages
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_community.tools import DuckDuckGoSearchRun
from langchain.tools import Calculator
from dotenv import load_dotenv


class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]


search_tool = DuckDuckGoSearchRun()
calculator = Calculator()
tools = [search_tool, calculator]

load_dotenv()
llm = ChatOpenAI("gpt-4o")
llm_with_tools = llm.bind_tools(tools)


def assistant(state: AgentState):
    system_prompt = """
    You are a well-educated research assistant with access to the web and a calculator. 
    Please answer the questions by outputting only the answer and nothing else.
    """

    system_message = SystemMessage(content=system_prompt)

    return {
        "messages": [llm_with_tools.invoke([system_message] + state["messages"])],
    }



class Agent:
    """
    A research assistant capable of searching the web and basic arithmetics.
    """

    def __init__(self):
        """
        Initializes the agent.
        """
        builder = StateGraph(AgentState)
        builder.add_node("assistant", assistant)
        builder.add_node("tools", ToolNode(tools))

        builder.add_edge(START, "assistant")
        builder.add_conditional_edges("assistant", tools_condition)

        builder.add_edge("tools", "assistant")
        self.agent = builder.compile()

    def __call__(self, question: str) -> str:
        """
        Answers a given question.

        Args: 
            question (str): Question to be answered.

        Returns:
            str: The answer to the question.
        """
        response = self.agent.invoke({"messages": [HumanMessage(content=f"Question:\n {question}")]})
        return response['messages'][-1].content