from typing import Annotated, TypedDict
from langchain_openai import ChatOpenAI
from tools import add, subtract, multiply, divide, exponentiate, web_search, paper_search, load_web_page, understand_image, transcribe_audio
from langchain_core.messages import AnyMessage
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.graph import START, StateGraph


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

class QAAgent:
    def __init__(self):
        print("QA Agent initialized.")
        self.agent = self.build_agent()

    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        response = self.agent.invoke({"messages": [HumanMessage(content=question)]})
        ret = response['messages'][-1].content
        print(f"Agent returning fixed answer: {ret}")
        return ret
    
    def build_agent(self):
        model = ChatOpenAI(model="gpt-4o-mini", temperature=0)
        tools = [add, subtract, multiply, divide, exponentiate, web_search, paper_search, load_web_page, understand_image, transcribe_audio]
        model_with_tools = model.bind_tools(tools)
        
        def assistant(state: AgentState):
            return {
                "messages": [model_with_tools.invoke(state["messages"])],
            }
        
        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")

        agent = builder.compile()
        return agent